This is an automated email from the ASF dual-hosted git repository. ueshin pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push: new 33e463c [SPARK-36769][PYTHON] Improve `filter` of single-indexed DataFrame 33e463c is described below commit 33e463ccf99d09ad8a743d32104f590e204da93d Author: Xinrong Meng <xinrong.m...@databricks.com> AuthorDate: Tue Sep 21 10:20:15 2021 -0700 [SPARK-36769][PYTHON] Improve `filter` of single-indexed DataFrame ### What changes were proposed in this pull request? Improve `filter` of single-indexed DataFrame by replacing a long Project with Filter or Join. ### Why are the changes needed? When the given `items` have too many elements, a long Project is introduced. We may replace that with `Column.isin` or joining depending on the length of `items` for better performance. ### Does this PR introduce _any_ user-facing change? No. ### How was this patch tested? Unit tests. Closes #33998 from xinrong-databricks/impr_filter. Authored-by: Xinrong Meng <xinrong.m...@databricks.com> Signed-off-by: Takuya UESHIN <ues...@databricks.com> --- python/pyspark/pandas/frame.py | 28 +++++++++++++++++++-------- python/pyspark/pandas/tests/test_dataframe.py | 7 +++++++ 2 files changed, 27 insertions(+), 8 deletions(-) diff --git a/python/pyspark/pandas/frame.py b/python/pyspark/pandas/frame.py index 09efef2..cba1db1 100644 --- a/python/pyspark/pandas/frame.py +++ b/python/pyspark/pandas/frame.py @@ -9995,13 +9995,25 @@ defaultdict(<class 'list'>, {'col..., 'col...})] raise ValueError("items should be a list-like object.") if axis == 0: if len(index_scols) == 1: - col = None - for item in items: - if col is None: - col = index_scols[0] == SF.lit(item) - else: - col = col | (index_scols[0] == SF.lit(item)) - elif len(index_scols) > 1: + if len(items) <= ps.get_option("compute.isin_limit"): + col = index_scols[0].isin([SF.lit(item) for item in items]) + return DataFrame(self._internal.with_filter(col)) + else: + item_sdf_col = verify_temp_column_name( + self._internal.spark_frame, "__item__" + ) + item_sdf = default_session().createDataFrame( + pd.DataFrame({item_sdf_col: items}) + ) + joined_sdf = self._internal.spark_frame.join( + other=F.broadcast(item_sdf), + on=(index_scols[0] == scol_for(item_sdf, item_sdf_col)), + how="semi", + ) + + return DataFrame(self._internal.with_new_sdf(joined_sdf)) + + else: # for multi-index col = None for item in items: @@ -10019,7 +10031,7 @@ defaultdict(<class 'list'>, {'col..., 'col...})] col = midx_col else: col = col | midx_col - return DataFrame(self._internal.with_filter(col)) + return DataFrame(self._internal.with_filter(col)) else: return self[items] elif like is not None: diff --git a/python/pyspark/pandas/tests/test_dataframe.py b/python/pyspark/pandas/tests/test_dataframe.py index 20aecc2..3cfbc03 100644 --- a/python/pyspark/pandas/tests/test_dataframe.py +++ b/python/pyspark/pandas/tests/test_dataframe.py @@ -4313,6 +4313,13 @@ class DataFrameTest(PandasOnSparkTestCase, SQLTestUtils): psdf.filter(items=["ab", "aa"], axis=0).sort_index(), pdf.filter(items=["ab", "aa"], axis=0).sort_index(), ) + + with option_context("compute.isin_limit", 0): + self.assert_eq( + psdf.filter(items=["ab", "aa"], axis=0).sort_index(), + pdf.filter(items=["ab", "aa"], axis=0).sort_index(), + ) + self.assert_eq( psdf.filter(items=["ba", "db"], axis=1).sort_index(), pdf.filter(items=["ba", "db"], axis=1).sort_index(), --------------------------------------------------------------------- To unsubscribe, e-mail: commits-unsubscr...@spark.apache.org For additional commands, e-mail: commits-h...@spark.apache.org