| |
| <!DOCTYPE html> |
| |
| <html> |
| <head> |
| <meta charset="utf-8" /> |
| <title>pyspark.testing.pandasutils — PySpark 3.5.3 documentation</title> |
| |
| <link href="../../../_static/styles/theme.css?digest=1999514e3f237ded88cf" rel="stylesheet"> |
| <link href="../../../_static/styles/pydata-sphinx-theme.css?digest=1999514e3f237ded88cf" rel="stylesheet"> |
| |
| |
| <link rel="stylesheet" |
| href="../../../_static/vendor/fontawesome/5.13.0/css/all.min.css"> |
| <link rel="preload" as="font" type="font/woff2" crossorigin |
| href="../../../_static/vendor/fontawesome/5.13.0/webfonts/fa-solid-900.woff2"> |
| <link rel="preload" as="font" type="font/woff2" crossorigin |
| href="../../../_static/vendor/fontawesome/5.13.0/webfonts/fa-brands-400.woff2"> |
| |
| |
| |
| |
| |
| <link rel="stylesheet" href="../../../_static/styles/pydata-sphinx-theme.css" type="text/css" /> |
| <link rel="stylesheet" href="../../../_static/pygments.css" type="text/css" /> |
| <link rel="stylesheet" type="text/css" href="../../../_static/copybutton.css" /> |
| <link rel="stylesheet" type="text/css" href="../../../_static/css/pyspark.css" /> |
| |
| <link rel="preload" as="script" href="../../../_static/scripts/pydata-sphinx-theme.js?digest=1999514e3f237ded88cf"> |
| |
| <script id="documentation_options" data-url_root="../../../" src="../../../_static/documentation_options.js"></script> |
| <script src="../../../_static/jquery.js"></script> |
| <script src="../../../_static/underscore.js"></script> |
| <script src="../../../_static/doctools.js"></script> |
| <script src="../../../_static/language_data.js"></script> |
| <script src="../../../_static/clipboard.min.js"></script> |
| <script src="../../../_static/copybutton.js"></script> |
| <script crossorigin="anonymous" integrity="sha256-Ae2Vz/4ePdIu6ZyI/5ZGsYnb+m0JlOmKPjt6XZ9JJkA=" src="https://cdnjs.cloudflare.com/ajax/libs/require.js/2.3.4/require.min.js"></script> |
| <script async="async" src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/latest.js?config=TeX-AMS-MML_HTMLorMML"></script> |
| <script type="text/x-mathjax-config">MathJax.Hub.Config({"tex2jax": {"inlineMath": [["$", "$"], ["\\(", "\\)"]], "processEscapes": true, "ignoreClass": "document", "processClass": "math|output_area"}})</script> |
| <link rel="canonical" href="https://spark.apache.org/docs/latest/api/python/_modules/pyspark/testing/pandasutils.html" /> |
| <link rel="search" title="Search" href="../../../search.html" /> |
| <meta name="viewport" content="width=device-width, initial-scale=1" /> |
| <meta name="docsearch:language" content="None"> |
| |
| |
| <!-- Google Analytics --> |
| |
| </head> |
| <body data-spy="scroll" data-target="#bd-toc-nav" data-offset="80"> |
| |
| <div class="container-fluid" id="banner"></div> |
| |
| |
| <nav class="navbar navbar-light navbar-expand-lg bg-light fixed-top bd-navbar" id="navbar-main"><div class="container-xl"> |
| |
| <div id="navbar-start"> |
| |
| |
| |
| <a class="navbar-brand" href="../../../index.html"> |
| <img src="../../../_static/spark-logo-reverse.png" class="logo" alt="logo"> |
| </a> |
| |
| |
| |
| </div> |
| |
| <button class="navbar-toggler" type="button" data-toggle="collapse" data-target="#navbar-collapsible" aria-controls="navbar-collapsible" aria-expanded="false" aria-label="Toggle navigation"> |
| <span class="navbar-toggler-icon"></span> |
| </button> |
| |
| |
| <div id="navbar-collapsible" class="col-lg-9 collapse navbar-collapse"> |
| <div id="navbar-center" class="mr-auto"> |
| |
| <div class="navbar-center-item"> |
| <ul id="navbar-main-elements" class="navbar-nav"> |
| <li class="toctree-l1 nav-item"> |
| <a class="reference internal nav-link" href="../../../index.html"> |
| Overview |
| </a> |
| </li> |
| |
| <li class="toctree-l1 nav-item"> |
| <a class="reference internal nav-link" href="../../../getting_started/index.html"> |
| Getting Started |
| </a> |
| </li> |
| |
| <li class="toctree-l1 nav-item"> |
| <a class="reference internal nav-link" href="../../../user_guide/index.html"> |
| User Guides |
| </a> |
| </li> |
| |
| <li class="toctree-l1 nav-item"> |
| <a class="reference internal nav-link" href="../../../reference/index.html"> |
| API Reference |
| </a> |
| </li> |
| |
| <li class="toctree-l1 nav-item"> |
| <a class="reference internal nav-link" href="../../../development/index.html"> |
| Development |
| </a> |
| </li> |
| |
| <li class="toctree-l1 nav-item"> |
| <a class="reference internal nav-link" href="../../../migration_guide/index.html"> |
| Migration Guides |
| </a> |
| </li> |
| |
| |
| </ul> |
| </div> |
| |
| </div> |
| |
| <div id="navbar-end"> |
| |
| <div class="navbar-end-item"> |
| <!-- |
| Licensed to the Apache Software Foundation (ASF) under one or more |
| contributor license agreements. See the NOTICE file distributed with |
| this work for additional information regarding copyright ownership. |
| The ASF licenses this file to You under the Apache License, Version 2.0 |
| (the "License"); you may not use this file except in compliance with |
| the License. You may obtain a copy of the License at |
| |
| http://www.apache.org/licenses/LICENSE-2.0 |
| |
| Unless required by applicable law or agreed to in writing, software |
| distributed under the License is distributed on an "AS IS" BASIS, |
| WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| See the License for the specific language governing permissions and |
| limitations under the License. |
| --> |
| |
| <div id="version-button" class="dropdown"> |
| <button type="button" class="btn btn-secondary btn-sm navbar-btn dropdown-toggle" id="version_switcher_button" data-toggle="dropdown"> |
| 3.5.3 |
| <span class="caret"></span> |
| </button> |
| <div id="version_switcher" class="dropdown-menu list-group-flush py-0" aria-labelledby="version_switcher_button"> |
| <!-- dropdown will be populated by javascript on page load --> |
| </div> |
| </div> |
| |
| <script type="text/javascript"> |
| // Function to construct the target URL from the JSON components |
| function buildURL(entry) { |
| var template = "https://spark.apache.org/docs/{version}/api/python/index.html"; // supplied by jinja |
| template = template.replace("{version}", entry.version); |
| return template; |
| } |
| |
| // Function to check if corresponding page path exists in other version of docs |
| // and, if so, go there instead of the homepage of the other docs version |
| function checkPageExistsAndRedirect(event) { |
| const currentFilePath = "_modules/pyspark/testing/pandasutils.html", |
| otherDocsHomepage = event.target.getAttribute("href"); |
| let tryUrl = `${otherDocsHomepage}${currentFilePath}`; |
| $.ajax({ |
| type: 'HEAD', |
| url: tryUrl, |
| // if the page exists, go there |
| success: function() { |
| location.href = tryUrl; |
| } |
| }).fail(function() { |
| location.href = otherDocsHomepage; |
| }); |
| return false; |
| } |
| |
| // Function to populate the version switcher |
| (function () { |
| // get JSON config |
| $.getJSON("https://spark.apache.org/static/versions.json", function(data, textStatus, jqXHR) { |
| // create the nodes first (before AJAX calls) to ensure the order is |
| // correct (for now, links will go to doc version homepage) |
| $.each(data, function(index, entry) { |
| // if no custom name specified (e.g., "latest"), use version string |
| if (!("name" in entry)) { |
| entry.name = entry.version; |
| } |
| // construct the appropriate URL, and add it to the dropdown |
| entry.url = buildURL(entry); |
| const node = document.createElement("a"); |
| node.setAttribute("class", "list-group-item list-group-item-action py-1"); |
| node.setAttribute("href", `${entry.url}`); |
| node.textContent = `${entry.name}`; |
| node.onclick = checkPageExistsAndRedirect; |
| $("#version_switcher").append(node); |
| }); |
| }); |
| })(); |
| </script> |
| </div> |
| |
| </div> |
| </div> |
| </div> |
| </nav> |
| |
| |
| <div class="container-xl"> |
| <div class="row"> |
| |
| |
| <!-- Only show if we have sidebars configured, else just a small margin --> |
| <div class="col-12 col-md-3 bd-sidebar"> |
| <div class="sidebar-start-items"><form class="bd-search d-flex align-items-center" action="../../../search.html" method="get"> |
| <i class="icon fas fa-search"></i> |
| <input type="search" class="form-control" name="q" id="search-input" placeholder="Search the docs ..." aria-label="Search the docs ..." autocomplete="off" > |
| </form><nav class="bd-links" id="bd-docs-nav" aria-label="Main navigation"> |
| <div class="bd-toc-item active"> |
| |
| </div> |
| </nav> |
| </div> |
| <div class="sidebar-end-items"> |
| </div> |
| </div> |
| |
| |
| |
| |
| <div class="d-none d-xl-block col-xl-2 bd-toc"> |
| |
| </div> |
| |
| |
| |
| |
| |
| |
| <main class="col-12 col-md-9 col-xl-7 py-md-5 pl-md-5 pr-md-4 bd-content" role="main"> |
| |
| <div> |
| |
| <h1>Source code for pyspark.testing.pandasutils</h1><div class="highlight"><pre> |
| <span></span><span class="c1">#</span> |
| <span class="c1"># Licensed to the Apache Software Foundation (ASF) under one or more</span> |
| <span class="c1"># contributor license agreements. See the NOTICE file distributed with</span> |
| <span class="c1"># this work for additional information regarding copyright ownership.</span> |
| <span class="c1"># The ASF licenses this file to You under the Apache License, Version 2.0</span> |
| <span class="c1"># (the "License"); you may not use this file except in compliance with</span> |
| <span class="c1"># the License. You may obtain a copy of the License at</span> |
| <span class="c1">#</span> |
| <span class="c1"># http://www.apache.org/licenses/LICENSE-2.0</span> |
| <span class="c1">#</span> |
| <span class="c1"># Unless required by applicable law or agreed to in writing, software</span> |
| <span class="c1"># distributed under the License is distributed on an "AS IS" BASIS,</span> |
| <span class="c1"># WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.</span> |
| <span class="c1"># See the License for the specific language governing permissions and</span> |
| <span class="c1"># limitations under the License.</span> |
| <span class="c1">#</span> |
| |
| <span class="kn">import</span> <span class="nn">functools</span> |
| <span class="kn">import</span> <span class="nn">shutil</span> |
| <span class="kn">import</span> <span class="nn">tempfile</span> |
| <span class="kn">import</span> <span class="nn">warnings</span> |
| <span class="kn">from</span> <span class="nn">contextlib</span> <span class="kn">import</span> <span class="n">contextmanager</span> |
| <span class="kn">from</span> <span class="nn">distutils.version</span> <span class="kn">import</span> <span class="n">LooseVersion</span> |
| <span class="kn">import</span> <span class="nn">decimal</span> |
| <span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Any</span><span class="p">,</span> <span class="n">Union</span><span class="p">,</span> <span class="n">TYPE_CHECKING</span> |
| |
| <span class="kn">import</span> <span class="nn">pyspark.pandas</span> <span class="k">as</span> <span class="nn">ps</span> |
| <span class="kn">from</span> <span class="nn">pyspark.pandas.frame</span> <span class="kn">import</span> <span class="n">DataFrame</span> |
| <span class="kn">from</span> <span class="nn">pyspark.pandas.indexes</span> <span class="kn">import</span> <span class="n">Index</span> |
| <span class="kn">from</span> <span class="nn">pyspark.pandas.series</span> <span class="kn">import</span> <span class="n">Series</span> |
| <span class="kn">from</span> <span class="nn">pyspark.pandas.utils</span> <span class="kn">import</span> <span class="n">SPARK_CONF_ARROW_ENABLED</span> |
| <span class="kn">from</span> <span class="nn">pyspark.testing.sqlutils</span> <span class="kn">import</span> <span class="n">ReusedSQLTestCase</span> |
| <span class="kn">from</span> <span class="nn">pyspark.errors</span> <span class="kn">import</span> <span class="n">PySparkAssertionError</span> |
| |
| <span class="n">tabulate_requirement_message</span> <span class="o">=</span> <span class="kc">None</span> |
| <span class="k">try</span><span class="p">:</span> |
| <span class="kn">from</span> <span class="nn">tabulate</span> <span class="kn">import</span> <span class="n">tabulate</span> |
| <span class="k">except</span> <span class="ne">ImportError</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span> |
| <span class="c1"># If tabulate requirement is not satisfied, skip related tests.</span> |
| <span class="n">tabulate_requirement_message</span> <span class="o">=</span> <span class="nb">str</span><span class="p">(</span><span class="n">e</span><span class="p">)</span> |
| <span class="n">have_tabulate</span> <span class="o">=</span> <span class="n">tabulate_requirement_message</span> <span class="ow">is</span> <span class="kc">None</span> |
| |
| <span class="n">matplotlib_requirement_message</span> <span class="o">=</span> <span class="kc">None</span> |
| <span class="k">try</span><span class="p">:</span> |
| <span class="kn">import</span> <span class="nn">matplotlib</span> |
| <span class="k">except</span> <span class="ne">ImportError</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span> |
| <span class="c1"># If matplotlib requirement is not satisfied, skip related tests.</span> |
| <span class="n">matplotlib_requirement_message</span> <span class="o">=</span> <span class="nb">str</span><span class="p">(</span><span class="n">e</span><span class="p">)</span> |
| <span class="n">have_matplotlib</span> <span class="o">=</span> <span class="n">matplotlib_requirement_message</span> <span class="ow">is</span> <span class="kc">None</span> |
| |
| <span class="n">plotly_requirement_message</span> <span class="o">=</span> <span class="kc">None</span> |
| <span class="k">try</span><span class="p">:</span> |
| <span class="kn">import</span> <span class="nn">plotly</span> |
| <span class="k">except</span> <span class="ne">ImportError</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span> |
| <span class="c1"># If plotly requirement is not satisfied, skip related tests.</span> |
| <span class="n">plotly_requirement_message</span> <span class="o">=</span> <span class="nb">str</span><span class="p">(</span><span class="n">e</span><span class="p">)</span> |
| <span class="n">have_plotly</span> <span class="o">=</span> <span class="n">plotly_requirement_message</span> <span class="ow">is</span> <span class="kc">None</span> |
| |
| <span class="k">try</span><span class="p">:</span> |
| <span class="kn">from</span> <span class="nn">pyspark.sql.pandas.utils</span> <span class="kn">import</span> <span class="n">require_minimum_pandas_version</span> |
| |
| <span class="n">require_minimum_pandas_version</span><span class="p">()</span> |
| <span class="kn">import</span> <span class="nn">pandas</span> <span class="k">as</span> <span class="nn">pd</span> |
| <span class="k">except</span> <span class="ne">ImportError</span><span class="p">:</span> |
| <span class="k">pass</span> |
| |
| <span class="n">__all__</span> <span class="o">=</span> <span class="p">[</span><span class="s2">"assertPandasOnSparkEqual"</span><span class="p">]</span> |
| |
| |
| <span class="k">def</span> <span class="nf">_assert_pandas_equal</span><span class="p">(</span> |
| <span class="n">left</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="n">pd</span><span class="o">.</span><span class="n">DataFrame</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Series</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Index</span><span class="p">],</span> |
| <span class="n">right</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="n">pd</span><span class="o">.</span><span class="n">DataFrame</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Series</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Index</span><span class="p">],</span> |
| <span class="n">checkExact</span><span class="p">:</span> <span class="nb">bool</span><span class="p">,</span> |
| <span class="p">):</span> |
| <span class="kn">from</span> <span class="nn">pandas.core.dtypes.common</span> <span class="kn">import</span> <span class="n">is_numeric_dtype</span> |
| <span class="kn">from</span> <span class="nn">pandas.testing</span> <span class="kn">import</span> <span class="n">assert_frame_equal</span><span class="p">,</span> <span class="n">assert_index_equal</span><span class="p">,</span> <span class="n">assert_series_equal</span> |
| |
| <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">left</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">DataFrame</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">right</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">DataFrame</span><span class="p">):</span> |
| <span class="k">try</span><span class="p">:</span> |
| <span class="k">if</span> <span class="n">LooseVersion</span><span class="p">(</span><span class="n">pd</span><span class="o">.</span><span class="n">__version__</span><span class="p">)</span> <span class="o">>=</span> <span class="n">LooseVersion</span><span class="p">(</span><span class="s2">"1.1"</span><span class="p">):</span> |
| <span class="n">kwargs</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">(</span><span class="n">check_freq</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span> |
| <span class="k">else</span><span class="p">:</span> |
| <span class="n">kwargs</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">()</span> |
| |
| <span class="k">if</span> <span class="n">LooseVersion</span><span class="p">(</span><span class="n">pd</span><span class="o">.</span><span class="n">__version__</span><span class="p">)</span> <span class="o"><</span> <span class="n">LooseVersion</span><span class="p">(</span><span class="s2">"1.1.1"</span><span class="p">):</span> |
| <span class="c1"># Due to https://github.com/pandas-dev/pandas/issues/35446</span> |
| <span class="n">checkExact</span> <span class="o">=</span> <span class="p">(</span> |
| <span class="n">checkExact</span> |
| <span class="ow">and</span> <span class="nb">all</span><span class="p">([</span><span class="n">is_numeric_dtype</span><span class="p">(</span><span class="n">dtype</span><span class="p">)</span> <span class="k">for</span> <span class="n">dtype</span> <span class="ow">in</span> <span class="n">left</span><span class="o">.</span><span class="n">dtypes</span><span class="p">])</span> |
| <span class="ow">and</span> <span class="nb">all</span><span class="p">([</span><span class="n">is_numeric_dtype</span><span class="p">(</span><span class="n">dtype</span><span class="p">)</span> <span class="k">for</span> <span class="n">dtype</span> <span class="ow">in</span> <span class="n">right</span><span class="o">.</span><span class="n">dtypes</span><span class="p">])</span> |
| <span class="p">)</span> |
| |
| <span class="n">assert_frame_equal</span><span class="p">(</span> |
| <span class="n">left</span><span class="p">,</span> |
| <span class="n">right</span><span class="p">,</span> |
| <span class="n">check_index_type</span><span class="o">=</span><span class="p">(</span><span class="s2">"equiv"</span> <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">index</span><span class="p">)</span> <span class="o">></span> <span class="mi">0</span> <span class="k">else</span> <span class="kc">False</span><span class="p">),</span> |
| <span class="n">check_column_type</span><span class="o">=</span><span class="p">(</span><span class="s2">"equiv"</span> <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">columns</span><span class="p">)</span> <span class="o">></span> <span class="mi">0</span> <span class="k">else</span> <span class="kc">False</span><span class="p">),</span> |
| <span class="n">check_exact</span><span class="o">=</span><span class="n">checkExact</span><span class="p">,</span> |
| <span class="o">**</span><span class="n">kwargs</span><span class="p">,</span> |
| <span class="p">)</span> |
| <span class="k">except</span> <span class="ne">AssertionError</span><span class="p">:</span> |
| <span class="k">raise</span> <span class="n">PySparkAssertionError</span><span class="p">(</span> |
| <span class="n">error_class</span><span class="o">=</span><span class="s2">"DIFFERENT_PANDAS_DATAFRAME"</span><span class="p">,</span> |
| <span class="n">message_parameters</span><span class="o">=</span><span class="p">{</span> |
| <span class="s2">"left"</span><span class="p">:</span> <span class="n">left</span><span class="o">.</span><span class="n">to_string</span><span class="p">(),</span> |
| <span class="s2">"left_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">dtypes</span><span class="p">),</span> |
| <span class="s2">"right"</span><span class="p">:</span> <span class="n">right</span><span class="o">.</span><span class="n">to_string</span><span class="p">(),</span> |
| <span class="s2">"right_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">right</span><span class="o">.</span><span class="n">dtypes</span><span class="p">),</span> |
| <span class="p">},</span> |
| <span class="p">)</span> |
| <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">left</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Series</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">right</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Series</span><span class="p">):</span> |
| <span class="k">try</span><span class="p">:</span> |
| <span class="k">if</span> <span class="n">LooseVersion</span><span class="p">(</span><span class="n">pd</span><span class="o">.</span><span class="n">__version__</span><span class="p">)</span> <span class="o">>=</span> <span class="n">LooseVersion</span><span class="p">(</span><span class="s2">"1.1"</span><span class="p">):</span> |
| <span class="n">kwargs</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">(</span><span class="n">check_freq</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span> |
| <span class="k">else</span><span class="p">:</span> |
| <span class="n">kwargs</span> <span class="o">=</span> <span class="nb">dict</span><span class="p">()</span> |
| <span class="k">if</span> <span class="n">LooseVersion</span><span class="p">(</span><span class="n">pd</span><span class="o">.</span><span class="n">__version__</span><span class="p">)</span> <span class="o"><</span> <span class="n">LooseVersion</span><span class="p">(</span><span class="s2">"1.1.1"</span><span class="p">):</span> |
| <span class="c1"># Due to https://github.com/pandas-dev/pandas/issues/35446</span> |
| <span class="n">checkExact</span> <span class="o">=</span> <span class="p">(</span> |
| <span class="n">checkExact</span> <span class="ow">and</span> <span class="n">is_numeric_dtype</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">dtype</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_numeric_dtype</span><span class="p">(</span><span class="n">right</span><span class="o">.</span><span class="n">dtype</span><span class="p">)</span> |
| <span class="p">)</span> |
| <span class="n">assert_series_equal</span><span class="p">(</span> |
| <span class="n">left</span><span class="p">,</span> |
| <span class="n">right</span><span class="p">,</span> |
| <span class="n">check_index_type</span><span class="o">=</span><span class="p">(</span><span class="s2">"equiv"</span> <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">index</span><span class="p">)</span> <span class="o">></span> <span class="mi">0</span> <span class="k">else</span> <span class="kc">False</span><span class="p">),</span> |
| <span class="n">check_exact</span><span class="o">=</span><span class="n">checkExact</span><span class="p">,</span> |
| <span class="o">**</span><span class="n">kwargs</span><span class="p">,</span> |
| <span class="p">)</span> |
| <span class="k">except</span> <span class="ne">AssertionError</span><span class="p">:</span> |
| <span class="k">raise</span> <span class="n">PySparkAssertionError</span><span class="p">(</span> |
| <span class="n">error_class</span><span class="o">=</span><span class="s2">"DIFFERENT_PANDAS_SERIES"</span><span class="p">,</span> |
| <span class="n">message_parameters</span><span class="o">=</span><span class="p">{</span> |
| <span class="s2">"left"</span><span class="p">:</span> <span class="n">left</span><span class="o">.</span><span class="n">to_string</span><span class="p">(),</span> |
| <span class="s2">"left_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">dtype</span><span class="p">),</span> |
| <span class="s2">"right"</span><span class="p">:</span> <span class="n">right</span><span class="o">.</span><span class="n">to_string</span><span class="p">(),</span> |
| <span class="s2">"right_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">right</span><span class="o">.</span><span class="n">dtype</span><span class="p">),</span> |
| <span class="p">},</span> |
| <span class="p">)</span> |
| <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">left</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Index</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">right</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Index</span><span class="p">):</span> |
| <span class="k">try</span><span class="p">:</span> |
| <span class="k">if</span> <span class="n">LooseVersion</span><span class="p">(</span><span class="n">pd</span><span class="o">.</span><span class="n">__version__</span><span class="p">)</span> <span class="o"><</span> <span class="n">LooseVersion</span><span class="p">(</span><span class="s2">"1.1.1"</span><span class="p">):</span> |
| <span class="c1"># Due to https://github.com/pandas-dev/pandas/issues/35446</span> |
| <span class="n">checkExact</span> <span class="o">=</span> <span class="p">(</span> |
| <span class="n">checkExact</span> <span class="ow">and</span> <span class="n">is_numeric_dtype</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">dtype</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_numeric_dtype</span><span class="p">(</span><span class="n">right</span><span class="o">.</span><span class="n">dtype</span><span class="p">)</span> |
| <span class="p">)</span> |
| <span class="n">assert_index_equal</span><span class="p">(</span><span class="n">left</span><span class="p">,</span> <span class="n">right</span><span class="p">,</span> <span class="n">check_exact</span><span class="o">=</span><span class="n">checkExact</span><span class="p">)</span> |
| <span class="k">except</span> <span class="ne">AssertionError</span><span class="p">:</span> |
| <span class="k">raise</span> <span class="n">PySparkAssertionError</span><span class="p">(</span> |
| <span class="n">error_class</span><span class="o">=</span><span class="s2">"DIFFERENT_PANDAS_INDEX"</span><span class="p">,</span> |
| <span class="n">message_parameters</span><span class="o">=</span><span class="p">{</span> |
| <span class="s2">"left"</span><span class="p">:</span> <span class="n">left</span><span class="p">,</span> |
| <span class="s2">"left_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">dtype</span><span class="p">),</span> |
| <span class="s2">"right"</span><span class="p">:</span> <span class="n">right</span><span class="p">,</span> |
| <span class="s2">"right_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">right</span><span class="o">.</span><span class="n">dtype</span><span class="p">),</span> |
| <span class="p">},</span> |
| <span class="p">)</span> |
| <span class="k">else</span><span class="p">:</span> |
| <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">"Unexpected values: (</span><span class="si">%s</span><span class="s2">, </span><span class="si">%s</span><span class="s2">)"</span> <span class="o">%</span> <span class="p">(</span><span class="n">left</span><span class="p">,</span> <span class="n">right</span><span class="p">))</span> |
| |
| |
| <span class="k">def</span> <span class="nf">_assert_pandas_almost_equal</span><span class="p">(</span> |
| <span class="n">left</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="n">pd</span><span class="o">.</span><span class="n">DataFrame</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Series</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Index</span><span class="p">],</span> |
| <span class="n">right</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="n">pd</span><span class="o">.</span><span class="n">DataFrame</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Series</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Index</span><span class="p">],</span> |
| <span class="n">rtol</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-5</span><span class="p">,</span> |
| <span class="n">atol</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-8</span><span class="p">,</span> |
| <span class="p">):</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> This function checks if given pandas objects approximately same,</span> |
| <span class="sd"> which means the conditions below:</span> |
| <span class="sd"> - Both objects are nullable</span> |
| <span class="sd"> - Compare decimals and floats, where two values a and b are approximately equal</span> |
| <span class="sd"> if they satisfy the following formula:</span> |
| <span class="sd"> absolute(a - b) <= (atol + rtol * absolute(b))</span> |
| <span class="sd"> where rtol=1e-5 and atol=1e-8 by default</span> |
| <span class="sd"> """</span> |
| |
| <span class="k">def</span> <span class="nf">compare_vals_approx</span><span class="p">(</span><span class="n">val1</span><span class="p">,</span> <span class="n">val2</span><span class="p">):</span> |
| <span class="c1"># compare vals for approximate equality</span> |
| <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">lval</span><span class="p">,</span> <span class="p">(</span><span class="nb">float</span><span class="p">,</span> <span class="n">decimal</span><span class="o">.</span><span class="n">Decimal</span><span class="p">))</span> <span class="ow">or</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">rval</span><span class="p">,</span> <span class="p">(</span><span class="nb">float</span><span class="p">,</span> <span class="n">decimal</span><span class="o">.</span><span class="n">Decimal</span><span class="p">)):</span> |
| <span class="k">if</span> <span class="nb">abs</span><span class="p">(</span><span class="nb">float</span><span class="p">(</span><span class="n">lval</span><span class="p">)</span> <span class="o">-</span> <span class="nb">float</span><span class="p">(</span><span class="n">rval</span><span class="p">))</span> <span class="o">></span> <span class="p">(</span><span class="n">atol</span> <span class="o">+</span> <span class="n">rtol</span> <span class="o">*</span> <span class="nb">abs</span><span class="p">(</span><span class="nb">float</span><span class="p">(</span><span class="n">rval</span><span class="p">))):</span> |
| <span class="k">return</span> <span class="kc">False</span> |
| <span class="k">elif</span> <span class="n">val1</span> <span class="o">!=</span> <span class="n">val2</span><span class="p">:</span> |
| <span class="k">return</span> <span class="kc">False</span> |
| <span class="k">return</span> <span class="kc">True</span> |
| |
| <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">left</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">DataFrame</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">right</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">DataFrame</span><span class="p">):</span> |
| <span class="k">if</span> <span class="n">left</span><span class="o">.</span><span class="n">shape</span> <span class="o">!=</span> <span class="n">right</span><span class="o">.</span><span class="n">shape</span><span class="p">:</span> |
| <span class="k">raise</span> <span class="n">PySparkAssertionError</span><span class="p">(</span> |
| <span class="n">error_class</span><span class="o">=</span><span class="s2">"DIFFERENT_PANDAS_DATAFRAME"</span><span class="p">,</span> |
| <span class="n">message_parameters</span><span class="o">=</span><span class="p">{</span> |
| <span class="s2">"left"</span><span class="p">:</span> <span class="n">left</span><span class="o">.</span><span class="n">to_string</span><span class="p">(),</span> |
| <span class="s2">"left_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">dtypes</span><span class="p">),</span> |
| <span class="s2">"right"</span><span class="p">:</span> <span class="n">right</span><span class="o">.</span><span class="n">to_string</span><span class="p">(),</span> |
| <span class="s2">"right_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">right</span><span class="o">.</span><span class="n">dtypes</span><span class="p">),</span> |
| <span class="p">},</span> |
| <span class="p">)</span> |
| <span class="k">for</span> <span class="n">lcol</span><span class="p">,</span> <span class="n">rcol</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">columns</span><span class="p">,</span> <span class="n">right</span><span class="o">.</span><span class="n">columns</span><span class="p">):</span> |
| <span class="k">if</span> <span class="n">lcol</span> <span class="o">!=</span> <span class="n">rcol</span><span class="p">:</span> |
| <span class="k">raise</span> <span class="n">PySparkAssertionError</span><span class="p">(</span> |
| <span class="n">error_class</span><span class="o">=</span><span class="s2">"DIFFERENT_PANDAS_DATAFRAME"</span><span class="p">,</span> |
| <span class="n">message_parameters</span><span class="o">=</span><span class="p">{</span> |
| <span class="s2">"left"</span><span class="p">:</span> <span class="n">left</span><span class="o">.</span><span class="n">to_string</span><span class="p">(),</span> |
| <span class="s2">"left_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">dtypes</span><span class="p">),</span> |
| <span class="s2">"right"</span><span class="p">:</span> <span class="n">right</span><span class="o">.</span><span class="n">to_string</span><span class="p">(),</span> |
| <span class="s2">"right_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">right</span><span class="o">.</span><span class="n">dtypes</span><span class="p">),</span> |
| <span class="p">},</span> |
| <span class="p">)</span> |
| <span class="k">for</span> <span class="n">lnull</span><span class="p">,</span> <span class="n">rnull</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">left</span><span class="p">[</span><span class="n">lcol</span><span class="p">]</span><span class="o">.</span><span class="n">isnull</span><span class="p">(),</span> <span class="n">right</span><span class="p">[</span><span class="n">rcol</span><span class="p">]</span><span class="o">.</span><span class="n">isnull</span><span class="p">()):</span> |
| <span class="k">if</span> <span class="n">lnull</span> <span class="o">!=</span> <span class="n">rnull</span><span class="p">:</span> |
| <span class="k">raise</span> <span class="n">PySparkAssertionError</span><span class="p">(</span> |
| <span class="n">error_class</span><span class="o">=</span><span class="s2">"DIFFERENT_PANDAS_DATAFRAME"</span><span class="p">,</span> |
| <span class="n">message_parameters</span><span class="o">=</span><span class="p">{</span> |
| <span class="s2">"left"</span><span class="p">:</span> <span class="n">left</span><span class="o">.</span><span class="n">to_string</span><span class="p">(),</span> |
| <span class="s2">"left_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">dtypes</span><span class="p">),</span> |
| <span class="s2">"right"</span><span class="p">:</span> <span class="n">right</span><span class="o">.</span><span class="n">to_string</span><span class="p">(),</span> |
| <span class="s2">"right_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">right</span><span class="o">.</span><span class="n">dtypes</span><span class="p">),</span> |
| <span class="p">},</span> |
| <span class="p">)</span> |
| <span class="k">for</span> <span class="n">lval</span><span class="p">,</span> <span class="n">rval</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">left</span><span class="p">[</span><span class="n">lcol</span><span class="p">]</span><span class="o">.</span><span class="n">dropna</span><span class="p">(),</span> <span class="n">right</span><span class="p">[</span><span class="n">rcol</span><span class="p">]</span><span class="o">.</span><span class="n">dropna</span><span class="p">()):</span> |
| <span class="k">if</span> <span class="ow">not</span> <span class="n">compare_vals_approx</span><span class="p">(</span><span class="n">lval</span><span class="p">,</span> <span class="n">rval</span><span class="p">):</span> |
| <span class="k">raise</span> <span class="n">PySparkAssertionError</span><span class="p">(</span> |
| <span class="n">error_class</span><span class="o">=</span><span class="s2">"DIFFERENT_PANDAS_DATAFRAME"</span><span class="p">,</span> |
| <span class="n">message_parameters</span><span class="o">=</span><span class="p">{</span> |
| <span class="s2">"left"</span><span class="p">:</span> <span class="n">left</span><span class="o">.</span><span class="n">to_string</span><span class="p">(),</span> |
| <span class="s2">"left_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">dtypes</span><span class="p">),</span> |
| <span class="s2">"right"</span><span class="p">:</span> <span class="n">right</span><span class="o">.</span><span class="n">to_string</span><span class="p">(),</span> |
| <span class="s2">"right_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">right</span><span class="o">.</span><span class="n">dtypes</span><span class="p">),</span> |
| <span class="p">},</span> |
| <span class="p">)</span> |
| <span class="k">if</span> <span class="n">left</span><span class="o">.</span><span class="n">columns</span><span class="o">.</span><span class="n">names</span> <span class="o">!=</span> <span class="n">right</span><span class="o">.</span><span class="n">columns</span><span class="o">.</span><span class="n">names</span><span class="p">:</span> |
| <span class="k">raise</span> <span class="n">PySparkAssertionError</span><span class="p">(</span> |
| <span class="n">error_class</span><span class="o">=</span><span class="s2">"DIFFERENT_PANDAS_DATAFRAME"</span><span class="p">,</span> |
| <span class="n">message_parameters</span><span class="o">=</span><span class="p">{</span> |
| <span class="s2">"left"</span><span class="p">:</span> <span class="n">left</span><span class="o">.</span><span class="n">to_string</span><span class="p">(),</span> |
| <span class="s2">"left_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">dtypes</span><span class="p">),</span> |
| <span class="s2">"right"</span><span class="p">:</span> <span class="n">right</span><span class="o">.</span><span class="n">to_string</span><span class="p">(),</span> |
| <span class="s2">"right_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">right</span><span class="o">.</span><span class="n">dtypes</span><span class="p">),</span> |
| <span class="p">},</span> |
| <span class="p">)</span> |
| <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">left</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Series</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">right</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Series</span><span class="p">):</span> |
| <span class="k">if</span> <span class="n">left</span><span class="o">.</span><span class="n">name</span> <span class="o">!=</span> <span class="n">right</span><span class="o">.</span><span class="n">name</span> <span class="ow">or</span> <span class="nb">len</span><span class="p">(</span><span class="n">left</span><span class="p">)</span> <span class="o">!=</span> <span class="nb">len</span><span class="p">(</span><span class="n">right</span><span class="p">):</span> |
| <span class="k">raise</span> <span class="n">PySparkAssertionError</span><span class="p">(</span> |
| <span class="n">error_class</span><span class="o">=</span><span class="s2">"DIFFERENT_PANDAS_SERIES"</span><span class="p">,</span> |
| <span class="n">message_parameters</span><span class="o">=</span><span class="p">{</span> |
| <span class="s2">"left"</span><span class="p">:</span> <span class="n">left</span><span class="o">.</span><span class="n">to_string</span><span class="p">(),</span> |
| <span class="s2">"left_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">dtype</span><span class="p">),</span> |
| <span class="s2">"right"</span><span class="p">:</span> <span class="n">right</span><span class="o">.</span><span class="n">to_string</span><span class="p">(),</span> |
| <span class="s2">"right_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">right</span><span class="o">.</span><span class="n">dtype</span><span class="p">),</span> |
| <span class="p">},</span> |
| <span class="p">)</span> |
| <span class="k">for</span> <span class="n">lnull</span><span class="p">,</span> <span class="n">rnull</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">isnull</span><span class="p">(),</span> <span class="n">right</span><span class="o">.</span><span class="n">isnull</span><span class="p">()):</span> |
| <span class="k">if</span> <span class="n">lnull</span> <span class="o">!=</span> <span class="n">rnull</span><span class="p">:</span> |
| <span class="k">raise</span> <span class="n">PySparkAssertionError</span><span class="p">(</span> |
| <span class="n">error_class</span><span class="o">=</span><span class="s2">"DIFFERENT_PANDAS_SERIES"</span><span class="p">,</span> |
| <span class="n">message_parameters</span><span class="o">=</span><span class="p">{</span> |
| <span class="s2">"left"</span><span class="p">:</span> <span class="n">left</span><span class="o">.</span><span class="n">to_string</span><span class="p">(),</span> |
| <span class="s2">"left_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">dtype</span><span class="p">),</span> |
| <span class="s2">"right"</span><span class="p">:</span> <span class="n">right</span><span class="o">.</span><span class="n">to_string</span><span class="p">(),</span> |
| <span class="s2">"right_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">right</span><span class="o">.</span><span class="n">dtype</span><span class="p">),</span> |
| <span class="p">},</span> |
| <span class="p">)</span> |
| <span class="k">for</span> <span class="n">lval</span><span class="p">,</span> <span class="n">rval</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">dropna</span><span class="p">(),</span> <span class="n">right</span><span class="o">.</span><span class="n">dropna</span><span class="p">()):</span> |
| <span class="k">if</span> <span class="ow">not</span> <span class="n">compare_vals_approx</span><span class="p">(</span><span class="n">lval</span><span class="p">,</span> <span class="n">rval</span><span class="p">):</span> |
| <span class="k">raise</span> <span class="n">PySparkAssertionError</span><span class="p">(</span> |
| <span class="n">error_class</span><span class="o">=</span><span class="s2">"DIFFERENT_PANDAS_SERIES"</span><span class="p">,</span> |
| <span class="n">message_parameters</span><span class="o">=</span><span class="p">{</span> |
| <span class="s2">"left"</span><span class="p">:</span> <span class="n">left</span><span class="o">.</span><span class="n">to_string</span><span class="p">(),</span> |
| <span class="s2">"left_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">dtype</span><span class="p">),</span> |
| <span class="s2">"right"</span><span class="p">:</span> <span class="n">right</span><span class="o">.</span><span class="n">to_string</span><span class="p">(),</span> |
| <span class="s2">"right_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">right</span><span class="o">.</span><span class="n">dtype</span><span class="p">),</span> |
| <span class="p">},</span> |
| <span class="p">)</span> |
| <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">left</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">MultiIndex</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">right</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">MultiIndex</span><span class="p">):</span> |
| <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">left</span><span class="p">)</span> <span class="o">!=</span> <span class="nb">len</span><span class="p">(</span><span class="n">right</span><span class="p">):</span> |
| <span class="k">raise</span> <span class="n">PySparkAssertionError</span><span class="p">(</span> |
| <span class="n">error_class</span><span class="o">=</span><span class="s2">"DIFFERENT_PANDAS_MULTIINDEX"</span><span class="p">,</span> |
| <span class="n">message_parameters</span><span class="o">=</span><span class="p">{</span> |
| <span class="s2">"left"</span><span class="p">:</span> <span class="n">left</span><span class="p">,</span> |
| <span class="s2">"left_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">dtype</span><span class="p">),</span> |
| <span class="s2">"right"</span><span class="p">:</span> <span class="n">right</span><span class="p">,</span> |
| <span class="s2">"right_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">right</span><span class="o">.</span><span class="n">dtype</span><span class="p">),</span> |
| <span class="p">},</span> |
| <span class="p">)</span> |
| <span class="k">for</span> <span class="n">lval</span><span class="p">,</span> <span class="n">rval</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">left</span><span class="p">,</span> <span class="n">right</span><span class="p">):</span> |
| <span class="k">if</span> <span class="ow">not</span> <span class="n">compare_vals_approx</span><span class="p">(</span><span class="n">lval</span><span class="p">,</span> <span class="n">rval</span><span class="p">):</span> |
| <span class="k">raise</span> <span class="n">PySparkAssertionError</span><span class="p">(</span> |
| <span class="n">error_class</span><span class="o">=</span><span class="s2">"DIFFERENT_PANDAS_MULTIINDEX"</span><span class="p">,</span> |
| <span class="n">message_parameters</span><span class="o">=</span><span class="p">{</span> |
| <span class="s2">"left"</span><span class="p">:</span> <span class="n">left</span><span class="p">,</span> |
| <span class="s2">"left_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">dtype</span><span class="p">),</span> |
| <span class="s2">"right"</span><span class="p">:</span> <span class="n">right</span><span class="p">,</span> |
| <span class="s2">"right_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">right</span><span class="o">.</span><span class="n">dtype</span><span class="p">),</span> |
| <span class="p">},</span> |
| <span class="p">)</span> |
| <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">left</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Index</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">right</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Index</span><span class="p">):</span> |
| <span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">left</span><span class="p">)</span> <span class="o">!=</span> <span class="nb">len</span><span class="p">(</span><span class="n">right</span><span class="p">):</span> |
| <span class="k">raise</span> <span class="n">PySparkAssertionError</span><span class="p">(</span> |
| <span class="n">error_class</span><span class="o">=</span><span class="s2">"DIFFERENT_PANDAS_INDEX"</span><span class="p">,</span> |
| <span class="n">message_parameters</span><span class="o">=</span><span class="p">{</span> |
| <span class="s2">"left"</span><span class="p">:</span> <span class="n">left</span><span class="p">,</span> |
| <span class="s2">"left_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">dtype</span><span class="p">),</span> |
| <span class="s2">"right"</span><span class="p">:</span> <span class="n">right</span><span class="p">,</span> |
| <span class="s2">"right_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">right</span><span class="o">.</span><span class="n">dtype</span><span class="p">),</span> |
| <span class="p">},</span> |
| <span class="p">)</span> |
| <span class="k">for</span> <span class="n">lnull</span><span class="p">,</span> <span class="n">rnull</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">isnull</span><span class="p">(),</span> <span class="n">right</span><span class="o">.</span><span class="n">isnull</span><span class="p">()):</span> |
| <span class="k">if</span> <span class="n">lnull</span> <span class="o">!=</span> <span class="n">rnull</span><span class="p">:</span> |
| <span class="k">raise</span> <span class="n">PySparkAssertionError</span><span class="p">(</span> |
| <span class="n">error_class</span><span class="o">=</span><span class="s2">"DIFFERENT_PANDAS_INDEX"</span><span class="p">,</span> |
| <span class="n">message_parameters</span><span class="o">=</span><span class="p">{</span> |
| <span class="s2">"left"</span><span class="p">:</span> <span class="n">left</span><span class="p">,</span> |
| <span class="s2">"left_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">dtype</span><span class="p">),</span> |
| <span class="s2">"right"</span><span class="p">:</span> <span class="n">right</span><span class="p">,</span> |
| <span class="s2">"right_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">right</span><span class="o">.</span><span class="n">dtype</span><span class="p">),</span> |
| <span class="p">},</span> |
| <span class="p">)</span> |
| <span class="k">for</span> <span class="n">lval</span><span class="p">,</span> <span class="n">rval</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">dropna</span><span class="p">(),</span> <span class="n">right</span><span class="o">.</span><span class="n">dropna</span><span class="p">()):</span> |
| <span class="k">if</span> <span class="ow">not</span> <span class="n">compare_vals_approx</span><span class="p">(</span><span class="n">lval</span><span class="p">,</span> <span class="n">rval</span><span class="p">):</span> |
| <span class="k">raise</span> <span class="n">PySparkAssertionError</span><span class="p">(</span> |
| <span class="n">error_class</span><span class="o">=</span><span class="s2">"DIFFERENT_PANDAS_INDEX"</span><span class="p">,</span> |
| <span class="n">message_parameters</span><span class="o">=</span><span class="p">{</span> |
| <span class="s2">"left"</span><span class="p">:</span> <span class="n">left</span><span class="p">,</span> |
| <span class="s2">"left_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">left</span><span class="o">.</span><span class="n">dtype</span><span class="p">),</span> |
| <span class="s2">"right"</span><span class="p">:</span> <span class="n">right</span><span class="p">,</span> |
| <span class="s2">"right_dtype"</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">right</span><span class="o">.</span><span class="n">dtype</span><span class="p">),</span> |
| <span class="p">},</span> |
| <span class="p">)</span> |
| <span class="k">else</span><span class="p">:</span> |
| <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">left</span><span class="p">,</span> <span class="p">(</span><span class="n">pd</span><span class="o">.</span><span class="n">DataFrame</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Series</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Index</span><span class="p">)):</span> |
| <span class="k">raise</span> <span class="n">PySparkAssertionError</span><span class="p">(</span> |
| <span class="n">error_class</span><span class="o">=</span><span class="s2">"INVALID_TYPE_DF_EQUALITY_ARG"</span><span class="p">,</span> |
| <span class="n">message_parameters</span><span class="o">=</span><span class="p">{</span> |
| <span class="s2">"expected_type"</span><span class="p">:</span> <span class="sa">f</span><span class="s2">"</span><span class="si">{</span><span class="n">pd</span><span class="o">.</span><span class="n">DataFrame</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">, "</span> |
| <span class="sa">f</span><span class="s2">"</span><span class="si">{</span><span class="n">pd</span><span class="o">.</span><span class="n">Series</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">, "</span> |
| <span class="sa">f</span><span class="s2">"</span><span class="si">{</span><span class="n">pd</span><span class="o">.</span><span class="n">Index</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">, "</span><span class="p">,</span> |
| <span class="s2">"arg_name"</span><span class="p">:</span> <span class="s2">"left"</span><span class="p">,</span> |
| <span class="s2">"actual_type"</span><span class="p">:</span> <span class="nb">type</span><span class="p">(</span><span class="n">left</span><span class="p">),</span> |
| <span class="p">},</span> |
| <span class="p">)</span> |
| <span class="k">elif</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">right</span><span class="p">,</span> <span class="p">(</span><span class="n">pd</span><span class="o">.</span><span class="n">DataFrame</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Series</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Index</span><span class="p">)):</span> |
| <span class="k">raise</span> <span class="n">PySparkAssertionError</span><span class="p">(</span> |
| <span class="n">error_class</span><span class="o">=</span><span class="s2">"INVALID_TYPE_DF_EQUALITY_ARG"</span><span class="p">,</span> |
| <span class="n">message_parameters</span><span class="o">=</span><span class="p">{</span> |
| <span class="s2">"expected_type"</span><span class="p">:</span> <span class="sa">f</span><span class="s2">"</span><span class="si">{</span><span class="n">pd</span><span class="o">.</span><span class="n">DataFrame</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">, "</span> |
| <span class="sa">f</span><span class="s2">"</span><span class="si">{</span><span class="n">pd</span><span class="o">.</span><span class="n">Series</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">, "</span> |
| <span class="sa">f</span><span class="s2">"</span><span class="si">{</span><span class="n">pd</span><span class="o">.</span><span class="n">Index</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">, "</span><span class="p">,</span> |
| <span class="s2">"arg_name"</span><span class="p">:</span> <span class="s2">"right"</span><span class="p">,</span> |
| <span class="s2">"actual_type"</span><span class="p">:</span> <span class="nb">type</span><span class="p">(</span><span class="n">right</span><span class="p">),</span> |
| <span class="p">},</span> |
| <span class="p">)</span> |
| |
| |
| <div class="viewcode-block" id="assertPandasOnSparkEqual"><a class="viewcode-back" href="../../../reference/api/pyspark.testing.assertPandasOnSparkEqual.html#pyspark.testing.assertPandasOnSparkEqual">[docs]</a><span class="k">def</span> <span class="nf">assertPandasOnSparkEqual</span><span class="p">(</span> |
| <span class="n">actual</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="n">DataFrame</span><span class="p">,</span> <span class="n">Series</span><span class="p">,</span> <span class="n">Index</span><span class="p">],</span> |
| <span class="n">expected</span><span class="p">:</span> <span class="n">Union</span><span class="p">[</span><span class="n">DataFrame</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">DataFrame</span><span class="p">,</span> <span class="n">Series</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Series</span><span class="p">,</span> <span class="n">Index</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Index</span><span class="p">],</span> |
| <span class="n">checkExact</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span> |
| <span class="n">almost</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span> |
| <span class="n">rtol</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-5</span><span class="p">,</span> |
| <span class="n">atol</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-8</span><span class="p">,</span> |
| <span class="n">checkRowOrder</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span> |
| <span class="p">):</span> |
| <span class="w"> </span><span class="sa">r</span><span class="sd">"""</span> |
| <span class="sd"> A util function to assert equality between actual (pandas-on-Spark object) and expected</span> |
| <span class="sd"> (pandas-on-Spark or pandas object).</span> |
| |
| <span class="sd"> .. versionadded:: 3.5.0</span> |
| |
| <span class="sd"> .. deprecated:: 3.5.1</span> |
| <span class="sd"> `assertPandasOnSparkEqual` will be removed in Spark 4.0.0.</span> |
| |
| <span class="sd"> Parameters</span> |
| <span class="sd"> ----------</span> |
| <span class="sd"> actual: pandas-on-Spark DataFrame, Series, or Index</span> |
| <span class="sd"> The object that is being compared or tested.</span> |
| <span class="sd"> expected: pandas-on-Spark or pandas DataFrame, Series, or Index</span> |
| <span class="sd"> The expected object, for comparison with the actual result.</span> |
| <span class="sd"> checkExact: bool, optional</span> |
| <span class="sd"> A flag indicating whether to compare exact equality.</span> |
| <span class="sd"> If set to 'True' (default), the data is compared exactly.</span> |
| <span class="sd"> If set to 'False', the data is compared less precisely, following pandas assert_frame_equal</span> |
| <span class="sd"> approximate comparison (see documentation for more details).</span> |
| <span class="sd"> almost: bool, optional</span> |
| <span class="sd"> A flag indicating whether to use unittest `assertAlmostEqual` or `assertEqual`.</span> |
| <span class="sd"> If set to 'True', the comparison is delegated to `unittest`'s `assertAlmostEqual`</span> |
| <span class="sd"> (see documentation for more details).</span> |
| <span class="sd"> If set to 'False' (default), the data is compared exactly with `unittest`'s</span> |
| <span class="sd"> `assertEqual`.</span> |
| <span class="sd"> rtol : float, optional</span> |
| <span class="sd"> The relative tolerance, used in asserting almost equality for float values in actual</span> |
| <span class="sd"> and expected. Set to 1e-5 by default. (See Notes)</span> |
| <span class="sd"> atol : float, optional</span> |
| <span class="sd"> The absolute tolerance, used in asserting almost equality for float values in actual</span> |
| <span class="sd"> and expected. Set to 1e-8 by default. (See Notes)</span> |
| <span class="sd"> checkRowOrder : bool, optional</span> |
| <span class="sd"> A flag indicating whether the order of rows should be considered in the comparison.</span> |
| <span class="sd"> If set to `False`, the row order is not taken into account.</span> |
| <span class="sd"> If set to `True` (default), the order of rows will be checked during comparison.</span> |
| <span class="sd"> (See Notes)</span> |
| |
| <span class="sd"> Notes</span> |
| <span class="sd"> -----</span> |
| <span class="sd"> For `checkRowOrder`, note that pandas-on-Spark DataFrame ordering is non-deterministic, unless</span> |
| <span class="sd"> explicitly sorted.</span> |
| |
| <span class="sd"> When `almost` is set to True, approximate equality will be asserted, where two values</span> |
| <span class="sd"> a and b are approximately equal if they satisfy the following formula:</span> |
| |
| <span class="sd"> ``absolute(a - b) <= (atol + rtol * absolute(b))``.</span> |
| |
| <span class="sd"> Examples</span> |
| <span class="sd"> --------</span> |
| <span class="sd"> >>> import pyspark.pandas as ps</span> |
| <span class="sd"> >>> psdf1 = ps.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6], 'c': [7, 8, 9]})</span> |
| <span class="sd"> >>> psdf2 = ps.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6], 'c': [7, 8, 9]})</span> |
| <span class="sd"> >>> assertPandasOnSparkEqual(psdf1, psdf2) # pass, ps.DataFrames are equal</span> |
| <span class="sd"> >>> s1 = ps.Series([212.32, 100.0001])</span> |
| <span class="sd"> >>> s2 = ps.Series([212.32, 100.0])</span> |
| <span class="sd"> >>> assertPandasOnSparkEqual(s1, s2, checkExact=False) # pass, ps.Series are approx equal</span> |
| <span class="sd"> >>> s1 = ps.Index([212.300001, 100.000])</span> |
| <span class="sd"> >>> s2 = ps.Index([212.3, 100.0001])</span> |
| <span class="sd"> >>> assertPandasOnSparkEqual(s1, s2, almost=True) # pass, ps.Index obj are almost equal</span> |
| <span class="sd"> """</span> |
| <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span> |
| <span class="s2">"`assertPandasOnSparkEqual` will be removed in Spark 4.0.0. "</span><span class="p">,</span> |
| <span class="ne">FutureWarning</span><span class="p">,</span> |
| <span class="p">)</span> |
| <span class="k">if</span> <span class="n">actual</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">expected</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span> |
| <span class="k">return</span> <span class="kc">True</span> |
| <span class="k">elif</span> <span class="n">actual</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="n">expected</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span> |
| <span class="k">return</span> <span class="kc">False</span> |
| |
| <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">actual</span><span class="p">,</span> <span class="p">(</span><span class="n">DataFrame</span><span class="p">,</span> <span class="n">Series</span><span class="p">,</span> <span class="n">Index</span><span class="p">)):</span> |
| <span class="k">raise</span> <span class="n">PySparkAssertionError</span><span class="p">(</span> |
| <span class="n">error_class</span><span class="o">=</span><span class="s2">"INVALID_TYPE_DF_EQUALITY_ARG"</span><span class="p">,</span> |
| <span class="n">message_parameters</span><span class="o">=</span><span class="p">{</span> |
| <span class="s2">"expected_type"</span><span class="p">:</span> <span class="sa">f</span><span class="s2">"</span><span class="si">{</span><span class="n">DataFrame</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">, </span><span class="si">{</span><span class="n">Series</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">, </span><span class="si">{</span><span class="n">Index</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">"</span><span class="p">,</span> |
| <span class="s2">"arg_name"</span><span class="p">:</span> <span class="s2">"actual"</span><span class="p">,</span> |
| <span class="s2">"actual_type"</span><span class="p">:</span> <span class="nb">type</span><span class="p">(</span><span class="n">actual</span><span class="p">),</span> |
| <span class="p">},</span> |
| <span class="p">)</span> |
| <span class="k">elif</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">expected</span><span class="p">,</span> <span class="p">(</span><span class="n">DataFrame</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">DataFrame</span><span class="p">,</span> <span class="n">Series</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Series</span><span class="p">,</span> <span class="n">Index</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Index</span><span class="p">)):</span> |
| <span class="k">raise</span> <span class="n">PySparkAssertionError</span><span class="p">(</span> |
| <span class="n">error_class</span><span class="o">=</span><span class="s2">"INVALID_TYPE_DF_EQUALITY_ARG"</span><span class="p">,</span> |
| <span class="n">message_parameters</span><span class="o">=</span><span class="p">{</span> |
| <span class="s2">"expected_type"</span><span class="p">:</span> <span class="sa">f</span><span class="s2">"</span><span class="si">{</span><span class="n">DataFrame</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">, "</span> |
| <span class="sa">f</span><span class="s2">"</span><span class="si">{</span><span class="n">pd</span><span class="o">.</span><span class="n">DataFrame</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">, "</span> |
| <span class="sa">f</span><span class="s2">"</span><span class="si">{</span><span class="n">Series</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">, "</span> |
| <span class="sa">f</span><span class="s2">"</span><span class="si">{</span><span class="n">pd</span><span class="o">.</span><span class="n">Series</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">, "</span> |
| <span class="sa">f</span><span class="s2">"</span><span class="si">{</span><span class="n">Index</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">"</span> |
| <span class="sa">f</span><span class="s2">"</span><span class="si">{</span><span class="n">pd</span><span class="o">.</span><span class="n">Index</span><span class="o">.</span><span class="vm">__name__</span><span class="si">}</span><span class="s2">, "</span><span class="p">,</span> |
| <span class="s2">"arg_name"</span><span class="p">:</span> <span class="s2">"expected"</span><span class="p">,</span> |
| <span class="s2">"actual_type"</span><span class="p">:</span> <span class="nb">type</span><span class="p">(</span><span class="n">expected</span><span class="p">),</span> |
| <span class="p">},</span> |
| <span class="p">)</span> |
| <span class="k">else</span><span class="p">:</span> |
| <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">actual</span><span class="p">,</span> <span class="p">(</span><span class="n">pd</span><span class="o">.</span><span class="n">DataFrame</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Index</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Series</span><span class="p">)):</span> |
| <span class="n">actual</span> <span class="o">=</span> <span class="n">actual</span><span class="o">.</span><span class="n">to_pandas</span><span class="p">()</span> |
| <span class="k">if</span> <span class="ow">not</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">expected</span><span class="p">,</span> <span class="p">(</span><span class="n">pd</span><span class="o">.</span><span class="n">DataFrame</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Index</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Series</span><span class="p">)):</span> |
| <span class="n">expected</span> <span class="o">=</span> <span class="n">expected</span><span class="o">.</span><span class="n">to_pandas</span><span class="p">()</span> |
| |
| <span class="k">if</span> <span class="ow">not</span> <span class="n">checkRowOrder</span><span class="p">:</span> |
| <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">actual</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">DataFrame</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">len</span><span class="p">(</span><span class="n">actual</span><span class="o">.</span><span class="n">columns</span><span class="p">)</span> <span class="o">></span> <span class="mi">0</span><span class="p">:</span> |
| <span class="n">actual</span> <span class="o">=</span> <span class="n">actual</span><span class="o">.</span><span class="n">sort_values</span><span class="p">(</span><span class="n">by</span><span class="o">=</span><span class="n">actual</span><span class="o">.</span><span class="n">columns</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">ignore_index</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span> |
| <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">expected</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">DataFrame</span><span class="p">)</span> <span class="ow">and</span> <span class="nb">len</span><span class="p">(</span><span class="n">expected</span><span class="o">.</span><span class="n">columns</span><span class="p">)</span> <span class="o">></span> <span class="mi">0</span><span class="p">:</span> |
| <span class="n">expected</span> <span class="o">=</span> <span class="n">expected</span><span class="o">.</span><span class="n">sort_values</span><span class="p">(</span><span class="n">by</span><span class="o">=</span><span class="n">expected</span><span class="o">.</span><span class="n">columns</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">ignore_index</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span> |
| |
| <span class="k">if</span> <span class="n">almost</span><span class="p">:</span> |
| <span class="n">_assert_pandas_almost_equal</span><span class="p">(</span><span class="n">actual</span><span class="p">,</span> <span class="n">expected</span><span class="p">,</span> <span class="n">rtol</span><span class="o">=</span><span class="n">rtol</span><span class="p">,</span> <span class="n">atol</span><span class="o">=</span><span class="n">atol</span><span class="p">)</span> |
| <span class="k">else</span><span class="p">:</span> |
| <span class="n">_assert_pandas_equal</span><span class="p">(</span><span class="n">actual</span><span class="p">,</span> <span class="n">expected</span><span class="p">,</span> <span class="n">checkExact</span><span class="o">=</span><span class="n">checkExact</span><span class="p">)</span></div> |
| |
| |
| <span class="k">class</span> <span class="nc">PandasOnSparkTestUtils</span><span class="p">:</span> |
| <span class="k">def</span> <span class="nf">convert_str_to_lambda</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">func</span><span class="p">:</span> <span class="nb">str</span><span class="p">):</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> This function converts `func` str to lambda call</span> |
| <span class="sd"> """</span> |
| <span class="k">return</span> <span class="k">lambda</span> <span class="n">x</span><span class="p">:</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">x</span><span class="p">,</span> <span class="n">func</span><span class="p">)()</span> |
| |
| <span class="k">def</span> <span class="nf">assertPandasEqual</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">left</span><span class="p">:</span> <span class="n">Any</span><span class="p">,</span> <span class="n">right</span><span class="p">:</span> <span class="n">Any</span><span class="p">,</span> <span class="n">check_exact</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">):</span> |
| <span class="n">_assert_pandas_equal</span><span class="p">(</span><span class="n">left</span><span class="p">,</span> <span class="n">right</span><span class="p">,</span> <span class="n">check_exact</span><span class="p">)</span> |
| |
| <span class="k">def</span> <span class="nf">assertPandasAlmostEqual</span><span class="p">(</span> |
| <span class="bp">self</span><span class="p">,</span> |
| <span class="n">left</span><span class="p">:</span> <span class="n">Any</span><span class="p">,</span> |
| <span class="n">right</span><span class="p">:</span> <span class="n">Any</span><span class="p">,</span> |
| <span class="n">rtol</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-5</span><span class="p">,</span> |
| <span class="n">atol</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-8</span><span class="p">,</span> |
| <span class="p">):</span> |
| <span class="n">_assert_pandas_almost_equal</span><span class="p">(</span><span class="n">left</span><span class="p">,</span> <span class="n">right</span><span class="p">,</span> <span class="n">rtol</span><span class="o">=</span><span class="n">rtol</span><span class="p">,</span> <span class="n">atol</span><span class="o">=</span><span class="n">atol</span><span class="p">)</span> |
| |
| <span class="k">def</span> <span class="nf">assert_eq</span><span class="p">(</span> |
| <span class="bp">self</span><span class="p">,</span> |
| <span class="n">left</span><span class="p">:</span> <span class="n">Any</span><span class="p">,</span> |
| <span class="n">right</span><span class="p">:</span> <span class="n">Any</span><span class="p">,</span> |
| <span class="n">check_exact</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span> |
| <span class="n">almost</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">,</span> |
| <span class="n">rtol</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-5</span><span class="p">,</span> |
| <span class="n">atol</span><span class="p">:</span> <span class="nb">float</span> <span class="o">=</span> <span class="mf">1e-8</span><span class="p">,</span> |
| <span class="n">check_row_order</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">True</span><span class="p">,</span> |
| <span class="p">):</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> Asserts if two arbitrary objects are equal or not. If given objects are Koalas DataFrame</span> |
| <span class="sd"> or Series, they are converted into pandas' and compared.</span> |
| |
| <span class="sd"> :param left: object to compare</span> |
| <span class="sd"> :param right: object to compare</span> |
| <span class="sd"> :param check_exact: if this is False, the comparison is done less precisely.</span> |
| <span class="sd"> :param almost: if this is enabled, the comparison asserts approximate equality</span> |
| <span class="sd"> for float and decimal values, where two values a and b are approximately equal</span> |
| <span class="sd"> if they satisfy the following formula:</span> |
| <span class="sd"> absolute(a - b) <= (atol + rtol * absolute(b))</span> |
| <span class="sd"> :param rtol: The relative tolerance, used in asserting approximate equality for</span> |
| <span class="sd"> float values. Set to 1e-5 by default.</span> |
| <span class="sd"> :param atol: The absolute tolerance, used in asserting approximate equality for</span> |
| <span class="sd"> float values in actual and expected. Set to 1e-8 by default.</span> |
| <span class="sd"> :param check_row_order: A flag indicating whether the order of rows should be considered</span> |
| <span class="sd"> in the comparison. If set to False, row order will be ignored.</span> |
| <span class="sd"> """</span> |
| <span class="kn">import</span> <span class="nn">pandas</span> <span class="k">as</span> <span class="nn">pd</span> |
| <span class="kn">from</span> <span class="nn">pandas.api.types</span> <span class="kn">import</span> <span class="n">is_list_like</span> |
| |
| <span class="c1"># for pandas-on-Spark DataFrames, allow choice to ignore row order</span> |
| <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">left</span><span class="p">,</span> <span class="p">(</span><span class="n">ps</span><span class="o">.</span><span class="n">DataFrame</span><span class="p">,</span> <span class="n">ps</span><span class="o">.</span><span class="n">Series</span><span class="p">,</span> <span class="n">ps</span><span class="o">.</span><span class="n">Index</span><span class="p">)):</span> |
| <span class="k">return</span> <span class="n">assertPandasOnSparkEqual</span><span class="p">(</span> |
| <span class="n">left</span><span class="p">,</span> |
| <span class="n">right</span><span class="p">,</span> |
| <span class="n">checkExact</span><span class="o">=</span><span class="n">check_exact</span><span class="p">,</span> |
| <span class="n">almost</span><span class="o">=</span><span class="n">almost</span><span class="p">,</span> |
| <span class="n">rtol</span><span class="o">=</span><span class="n">rtol</span><span class="p">,</span> |
| <span class="n">atol</span><span class="o">=</span><span class="n">atol</span><span class="p">,</span> |
| <span class="n">checkRowOrder</span><span class="o">=</span><span class="n">check_row_order</span><span class="p">,</span> |
| <span class="p">)</span> |
| |
| <span class="n">lobj</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_to_pandas</span><span class="p">(</span><span class="n">left</span><span class="p">)</span> |
| <span class="n">robj</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_to_pandas</span><span class="p">(</span><span class="n">right</span><span class="p">)</span> |
| <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">lobj</span><span class="p">,</span> <span class="p">(</span><span class="n">pd</span><span class="o">.</span><span class="n">DataFrame</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Series</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Index</span><span class="p">)):</span> |
| <span class="k">if</span> <span class="n">almost</span><span class="p">:</span> |
| <span class="n">_assert_pandas_almost_equal</span><span class="p">(</span><span class="n">lobj</span><span class="p">,</span> <span class="n">robj</span><span class="p">,</span> <span class="n">rtol</span><span class="o">=</span><span class="n">rtol</span><span class="p">,</span> <span class="n">atol</span><span class="o">=</span><span class="n">atol</span><span class="p">)</span> |
| <span class="k">else</span><span class="p">:</span> |
| <span class="n">_assert_pandas_equal</span><span class="p">(</span><span class="n">lobj</span><span class="p">,</span> <span class="n">robj</span><span class="p">,</span> <span class="n">checkExact</span><span class="o">=</span><span class="n">check_exact</span><span class="p">)</span> |
| <span class="k">elif</span> <span class="n">is_list_like</span><span class="p">(</span><span class="n">lobj</span><span class="p">)</span> <span class="ow">and</span> <span class="n">is_list_like</span><span class="p">(</span><span class="n">robj</span><span class="p">):</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">assertTrue</span><span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">left</span><span class="p">)</span> <span class="o">==</span> <span class="nb">len</span><span class="p">(</span><span class="n">right</span><span class="p">))</span> |
| <span class="k">for</span> <span class="n">litem</span><span class="p">,</span> <span class="n">ritem</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">left</span><span class="p">,</span> <span class="n">right</span><span class="p">):</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">assert_eq</span><span class="p">(</span><span class="n">litem</span><span class="p">,</span> <span class="n">ritem</span><span class="p">,</span> <span class="n">check_exact</span><span class="o">=</span><span class="n">check_exact</span><span class="p">,</span> <span class="n">almost</span><span class="o">=</span><span class="n">almost</span><span class="p">)</span> |
| <span class="k">elif</span> <span class="p">(</span><span class="n">lobj</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">pd</span><span class="o">.</span><span class="n">isna</span><span class="p">(</span><span class="n">lobj</span><span class="p">))</span> <span class="ow">and</span> <span class="p">(</span><span class="n">robj</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="n">pd</span><span class="o">.</span><span class="n">isna</span><span class="p">(</span><span class="n">robj</span><span class="p">)):</span> |
| <span class="k">pass</span> |
| <span class="k">else</span><span class="p">:</span> |
| <span class="k">if</span> <span class="n">almost</span><span class="p">:</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">assertAlmostEqual</span><span class="p">(</span><span class="n">lobj</span><span class="p">,</span> <span class="n">robj</span><span class="p">)</span> |
| <span class="k">else</span><span class="p">:</span> |
| <span class="bp">self</span><span class="o">.</span><span class="n">assertEqual</span><span class="p">(</span><span class="n">lobj</span><span class="p">,</span> <span class="n">robj</span><span class="p">)</span> |
| |
| <span class="nd">@staticmethod</span> |
| <span class="k">def</span> <span class="nf">_to_pandas</span><span class="p">(</span><span class="n">obj</span><span class="p">:</span> <span class="n">Any</span><span class="p">):</span> |
| <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">obj</span><span class="p">,</span> <span class="p">(</span><span class="n">DataFrame</span><span class="p">,</span> <span class="n">Series</span><span class="p">,</span> <span class="n">Index</span><span class="p">)):</span> |
| <span class="k">return</span> <span class="n">obj</span><span class="o">.</span><span class="n">to_pandas</span><span class="p">()</span> |
| <span class="k">else</span><span class="p">:</span> |
| <span class="k">return</span> <span class="n">obj</span> |
| |
| |
| <span class="k">class</span> <span class="nc">PandasOnSparkTestCase</span><span class="p">(</span><span class="n">ReusedSQLTestCase</span><span class="p">,</span> <span class="n">PandasOnSparkTestUtils</span><span class="p">):</span> |
| <span class="nd">@classmethod</span> |
| <span class="k">def</span> <span class="nf">setUpClass</span><span class="p">(</span><span class="bp">cls</span><span class="p">):</span> |
| <span class="nb">super</span><span class="p">(</span><span class="n">PandasOnSparkTestCase</span><span class="p">,</span> <span class="bp">cls</span><span class="p">)</span><span class="o">.</span><span class="n">setUpClass</span><span class="p">()</span> |
| <span class="bp">cls</span><span class="o">.</span><span class="n">spark</span><span class="o">.</span><span class="n">conf</span><span class="o">.</span><span class="n">set</span><span class="p">(</span><span class="n">SPARK_CONF_ARROW_ENABLED</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span> |
| |
| |
| <span class="k">class</span> <span class="nc">TestUtils</span><span class="p">:</span> |
| <span class="nd">@contextmanager</span> |
| <span class="k">def</span> <span class="nf">temp_dir</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span> |
| <span class="n">tmp</span> <span class="o">=</span> <span class="n">tempfile</span><span class="o">.</span><span class="n">mkdtemp</span><span class="p">()</span> |
| <span class="k">try</span><span class="p">:</span> |
| <span class="k">yield</span> <span class="n">tmp</span> |
| <span class="k">finally</span><span class="p">:</span> |
| <span class="n">shutil</span><span class="o">.</span><span class="n">rmtree</span><span class="p">(</span><span class="n">tmp</span><span class="p">)</span> |
| |
| <span class="nd">@contextmanager</span> |
| <span class="k">def</span> <span class="nf">temp_file</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span> |
| <span class="k">with</span> <span class="bp">self</span><span class="o">.</span><span class="n">temp_dir</span><span class="p">()</span> <span class="k">as</span> <span class="n">tmp</span><span class="p">:</span> |
| <span class="k">yield</span> <span class="n">tempfile</span><span class="o">.</span><span class="n">mkstemp</span><span class="p">(</span><span class="nb">dir</span><span class="o">=</span><span class="n">tmp</span><span class="p">)[</span><span class="mi">1</span><span class="p">]</span> |
| |
| |
| <span class="k">class</span> <span class="nc">ComparisonTestBase</span><span class="p">(</span><span class="n">PandasOnSparkTestCase</span><span class="p">):</span> |
| <span class="nd">@property</span> |
| <span class="k">def</span> <span class="nf">psdf</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span> |
| <span class="k">return</span> <span class="n">ps</span><span class="o">.</span><span class="n">from_pandas</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">pdf</span><span class="p">)</span> |
| |
| <span class="nd">@property</span> |
| <span class="k">def</span> <span class="nf">pdf</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span> |
| <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">psdf</span><span class="o">.</span><span class="n">to_pandas</span><span class="p">()</span> |
| |
| |
| <span class="k">def</span> <span class="nf">compare_both</span><span class="p">(</span><span class="n">f</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">almost</span><span class="o">=</span><span class="kc">True</span><span class="p">):</span> |
| |
| <span class="k">if</span> <span class="n">f</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span> |
| <span class="k">return</span> <span class="n">functools</span><span class="o">.</span><span class="n">partial</span><span class="p">(</span><span class="n">compare_both</span><span class="p">,</span> <span class="n">almost</span><span class="o">=</span><span class="n">almost</span><span class="p">)</span> |
| <span class="k">elif</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">f</span><span class="p">,</span> <span class="nb">bool</span><span class="p">):</span> |
| <span class="k">return</span> <span class="n">functools</span><span class="o">.</span><span class="n">partial</span><span class="p">(</span><span class="n">compare_both</span><span class="p">,</span> <span class="n">almost</span><span class="o">=</span><span class="n">f</span><span class="p">)</span> |
| |
| <span class="nd">@functools</span><span class="o">.</span><span class="n">wraps</span><span class="p">(</span><span class="n">f</span><span class="p">)</span> |
| <span class="k">def</span> <span class="nf">wrapped</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span> |
| <span class="k">if</span> <span class="n">almost</span><span class="p">:</span> |
| <span class="n">compare</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">assertPandasAlmostEqual</span> |
| <span class="k">else</span><span class="p">:</span> |
| <span class="n">compare</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">assertPandasEqual</span> |
| |
| <span class="k">for</span> <span class="n">result_pandas</span><span class="p">,</span> <span class="n">result_spark</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">f</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">pdf</span><span class="p">),</span> <span class="n">f</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">psdf</span><span class="p">)):</span> |
| <span class="n">compare</span><span class="p">(</span><span class="n">result_pandas</span><span class="p">,</span> <span class="n">result_spark</span><span class="o">.</span><span class="n">to_pandas</span><span class="p">())</span> |
| |
| <span class="k">return</span> <span class="n">wrapped</span> |
| |
| |
| <span class="nd">@contextmanager</span> |
| <span class="k">def</span> <span class="nf">assert_produces_warning</span><span class="p">(</span> |
| <span class="n">expected_warning</span><span class="o">=</span><span class="ne">Warning</span><span class="p">,</span> |
| <span class="n">filter_level</span><span class="o">=</span><span class="s2">"always"</span><span class="p">,</span> |
| <span class="n">check_stacklevel</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> |
| <span class="n">raise_on_extra_warnings</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> |
| <span class="p">):</span> |
| <span class="w"> </span><span class="sd">"""</span> |
| <span class="sd"> Context manager for running code expected to either raise a specific</span> |
| <span class="sd"> warning, or not raise any warnings. Verifies that the code raises the</span> |
| <span class="sd"> expected warning, and that it does not raise any other unexpected</span> |
| <span class="sd"> warnings. It is basically a wrapper around ``warnings.catch_warnings``.</span> |
| |
| <span class="sd"> Notes</span> |
| <span class="sd"> -----</span> |
| <span class="sd"> Replicated from pandas/_testing/_warnings.py.</span> |
| |
| <span class="sd"> Parameters</span> |
| <span class="sd"> ----------</span> |
| <span class="sd"> expected_warning : {Warning, False, None}, default Warning</span> |
| <span class="sd"> The type of Exception raised. ``exception.Warning`` is the base</span> |
| <span class="sd"> class for all warnings. To check that no warning is returned,</span> |
| <span class="sd"> specify ``False`` or ``None``.</span> |
| <span class="sd"> filter_level : str or None, default "always"</span> |
| <span class="sd"> Specifies whether warnings are ignored, displayed, or turned</span> |
| <span class="sd"> into errors.</span> |
| <span class="sd"> Valid values are:</span> |
| <span class="sd"> * "error" - turns matching warnings into exceptions</span> |
| <span class="sd"> * "ignore" - discard the warning</span> |
| <span class="sd"> * "always" - always emit a warning</span> |
| <span class="sd"> * "default" - print the warning the first time it is generated</span> |
| <span class="sd"> from each location</span> |
| <span class="sd"> * "module" - print the warning the first time it is generated</span> |
| <span class="sd"> from each module</span> |
| <span class="sd"> * "once" - print the warning the first time it is generated</span> |
| <span class="sd"> check_stacklevel : bool, default True</span> |
| <span class="sd"> If True, displays the line that called the function containing</span> |
| <span class="sd"> the warning to show were the function is called. Otherwise, the</span> |
| <span class="sd"> line that implements the function is displayed.</span> |
| <span class="sd"> raise_on_extra_warnings : bool, default True</span> |
| <span class="sd"> Whether extra warnings not of the type `expected_warning` should</span> |
| <span class="sd"> cause the test to fail.</span> |
| |
| <span class="sd"> Examples</span> |
| <span class="sd"> --------</span> |
| <span class="sd"> >>> import warnings</span> |
| <span class="sd"> >>> with assert_produces_warning():</span> |
| <span class="sd"> ... warnings.warn(UserWarning())</span> |
| <span class="sd"> ...</span> |
| <span class="sd"> >>> with assert_produces_warning(False): # doctest: +SKIP</span> |
| <span class="sd"> ... warnings.warn(RuntimeWarning())</span> |
| <span class="sd"> ...</span> |
| <span class="sd"> Traceback (most recent call last):</span> |
| <span class="sd"> ...</span> |
| <span class="sd"> AssertionError: Caused unexpected warning(s): ['RuntimeWarning'].</span> |
| <span class="sd"> >>> with assert_produces_warning(UserWarning): # doctest: +SKIP</span> |
| <span class="sd"> ... warnings.warn(RuntimeWarning())</span> |
| <span class="sd"> Traceback (most recent call last):</span> |
| <span class="sd"> ...</span> |
| <span class="sd"> AssertionError: Did not see expected warning of class 'UserWarning'</span> |
| <span class="sd"> ..warn:: This is *not* thread-safe.</span> |
| <span class="sd"> """</span> |
| <span class="n">__tracebackhide__</span> <span class="o">=</span> <span class="kc">True</span> |
| |
| <span class="k">with</span> <span class="n">warnings</span><span class="o">.</span><span class="n">catch_warnings</span><span class="p">(</span><span class="n">record</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span> <span class="k">as</span> <span class="n">w</span><span class="p">:</span> |
| |
| <span class="n">saw_warning</span> <span class="o">=</span> <span class="kc">False</span> |
| <span class="n">warnings</span><span class="o">.</span><span class="n">simplefilter</span><span class="p">(</span><span class="n">filter_level</span><span class="p">)</span> |
| <span class="k">yield</span> <span class="n">w</span> |
| <span class="n">extra_warnings</span> <span class="o">=</span> <span class="p">[]</span> |
| |
| <span class="k">for</span> <span class="n">actual_warning</span> <span class="ow">in</span> <span class="n">w</span><span class="p">:</span> |
| <span class="k">if</span> <span class="n">expected_warning</span> <span class="ow">and</span> <span class="nb">issubclass</span><span class="p">(</span><span class="n">actual_warning</span><span class="o">.</span><span class="n">category</span><span class="p">,</span> <span class="n">expected_warning</span><span class="p">):</span> |
| <span class="n">saw_warning</span> <span class="o">=</span> <span class="kc">True</span> |
| |
| <span class="k">if</span> <span class="n">check_stacklevel</span> <span class="ow">and</span> <span class="nb">issubclass</span><span class="p">(</span> |
| <span class="n">actual_warning</span><span class="o">.</span><span class="n">category</span><span class="p">,</span> <span class="p">(</span><span class="ne">FutureWarning</span><span class="p">,</span> <span class="ne">DeprecationWarning</span><span class="p">)</span> |
| <span class="p">):</span> |
| <span class="kn">from</span> <span class="nn">inspect</span> <span class="kn">import</span> <span class="n">getframeinfo</span><span class="p">,</span> <span class="n">stack</span> |
| |
| <span class="n">caller</span> <span class="o">=</span> <span class="n">getframeinfo</span><span class="p">(</span><span class="n">stack</span><span class="p">()[</span><span class="mi">2</span><span class="p">][</span><span class="mi">0</span><span class="p">])</span> |
| <span class="n">msg</span> <span class="o">=</span> <span class="p">(</span> |
| <span class="s2">"Warning not set with correct stacklevel. "</span><span class="p">,</span> |
| <span class="s2">"File where warning is raised: </span><span class="si">{}</span><span class="s2"> != "</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">actual_warning</span><span class="o">.</span><span class="n">filename</span><span class="p">),</span> |
| <span class="s2">"</span><span class="si">{}</span><span class="s2">. Warning message: </span><span class="si">{}</span><span class="s2">"</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="n">caller</span><span class="o">.</span><span class="n">filename</span><span class="p">,</span> <span class="n">actual_warning</span><span class="o">.</span><span class="n">message</span><span class="p">),</span> |
| <span class="p">)</span> |
| <span class="k">assert</span> <span class="n">actual_warning</span><span class="o">.</span><span class="n">filename</span> <span class="o">==</span> <span class="n">caller</span><span class="o">.</span><span class="n">filename</span><span class="p">,</span> <span class="n">msg</span> |
| <span class="k">else</span><span class="p">:</span> |
| <span class="n">extra_warnings</span><span class="o">.</span><span class="n">append</span><span class="p">(</span> |
| <span class="p">(</span> |
| <span class="n">actual_warning</span><span class="o">.</span><span class="n">category</span><span class="o">.</span><span class="vm">__name__</span><span class="p">,</span> |
| <span class="n">actual_warning</span><span class="o">.</span><span class="n">message</span><span class="p">,</span> |
| <span class="n">actual_warning</span><span class="o">.</span><span class="n">filename</span><span class="p">,</span> |
| <span class="n">actual_warning</span><span class="o">.</span><span class="n">lineno</span><span class="p">,</span> |
| <span class="p">)</span> |
| <span class="p">)</span> |
| <span class="k">if</span> <span class="n">expected_warning</span><span class="p">:</span> |
| <span class="n">msg</span> <span class="o">=</span> <span class="s2">"Did not see expected warning of class </span><span class="si">{}</span><span class="s2">"</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="nb">repr</span><span class="p">(</span><span class="n">expected_warning</span><span class="o">.</span><span class="vm">__name__</span><span class="p">))</span> |
| <span class="k">assert</span> <span class="n">saw_warning</span><span class="p">,</span> <span class="n">msg</span> |
| <span class="k">if</span> <span class="n">raise_on_extra_warnings</span> <span class="ow">and</span> <span class="n">extra_warnings</span><span class="p">:</span> |
| <span class="k">raise</span> <span class="ne">AssertionError</span><span class="p">(</span><span class="s2">"Caused unexpected warning(s): </span><span class="si">{}</span><span class="s2">"</span><span class="o">.</span><span class="n">format</span><span class="p">(</span><span class="nb">repr</span><span class="p">(</span><span class="n">extra_warnings</span><span class="p">)))</span> |
| </pre></div> |
| |
| </div> |
| |
| |
| <!-- Previous / next buttons --> |
| <div class='prev-next-area'> |
| </div> |
| |
| </main> |
| |
| |
| </div> |
| </div> |
| |
| <script src="../../../_static/scripts/pydata-sphinx-theme.js?digest=1999514e3f237ded88cf"></script> |
| <footer class="footer mt-5 mt-md-0"> |
| <div class="container"> |
| |
| <div class="footer-item"> |
| <p class="copyright"> |
| © Copyright .<br> |
| </p> |
| </div> |
| |
| <div class="footer-item"> |
| <p class="sphinx-version"> |
| Created using <a href="http://sphinx-doc.org/">Sphinx</a> 3.0.4.<br> |
| </p> |
| </div> |
| |
| </div> |
| </footer> |
| </body> |
| </html> |
