From 6544fc9ce7a55325385419d3d9777a9733d2c068 Mon Sep 17 00:00:00 2001 From: yassun7010 Date: Thu, 15 Feb 2024 21:59:13 +0900 Subject: [PATCH] feat: change limit behavor. --- .../src/turu/snowflake/record/async_record_cursor.py | 7 ++----- turu-snowflake/src/turu/snowflake/record/record_cursor.py | 7 ++----- turu-snowflake/tests/turu/test_snowflake.py | 2 +- turu-snowflake/tests/turu/test_snowflake_async.py | 2 +- 4 files changed, 6 insertions(+), 12 deletions(-) diff --git a/turu-snowflake/src/turu/snowflake/record/async_record_cursor.py b/turu-snowflake/src/turu/snowflake/record/async_record_cursor.py index 641751b..5f9b4e5 100644 --- a/turu-snowflake/src/turu/snowflake/record/async_record_cursor.py +++ b/turu-snowflake/src/turu/snowflake/record/async_record_cursor.py @@ -8,12 +8,9 @@ async def fetch_pandas_all(self, **kwargs) -> "PandasDataFrame": if isinstance(self._recorder, turu.core.record.CsvRecorder): if limit := self._recorder._options.get("limit"): - record_df = df.head(limit) + df = df.head(limit) - else: - record_df = df - - record_df.to_csv( + df.to_csv( self._recorder.file, index=False, header=self._recorder._options.get("header", True), diff --git a/turu-snowflake/src/turu/snowflake/record/record_cursor.py b/turu-snowflake/src/turu/snowflake/record/record_cursor.py index 2602011..1dc6b86 100644 --- a/turu-snowflake/src/turu/snowflake/record/record_cursor.py +++ b/turu-snowflake/src/turu/snowflake/record/record_cursor.py @@ -8,12 +8,9 @@ def fetch_pandas_all(self, **kwargs) -> "PandasDataFrame": if isinstance(self._recorder, turu.core.record.CsvRecorder): if limit := self._recorder._options.get("limit"): - record_df = df.head(limit) + df = df.head(limit) - else: - record_df = df - - record_df.to_csv( + df.to_csv( self._recorder.file, index=False, header=self._recorder._options.get("header", True), diff --git a/turu-snowflake/tests/turu/test_snowflake.py b/turu-snowflake/tests/turu/test_snowflake.py index 5fd4bc5..f91f7d0 100644 --- a/turu-snowflake/tests/turu/test_snowflake.py +++ b/turu-snowflake/tests/turu/test_snowflake.py @@ -311,7 +311,7 @@ def test_record_pandas_dataframe_with_limit_option( ), limit=2, ) as cursor: - expected = pd.DataFrame({"ID": list(range(1, 10))}, dtype="object") + expected = pd.DataFrame({"ID": [1, 2]}, dtype="object") assert_frame_equal(cursor.fetch_pandas_all(), expected) diff --git a/turu-snowflake/tests/turu/test_snowflake_async.py b/turu-snowflake/tests/turu/test_snowflake_async.py index 692e725..ec6378e 100644 --- a/turu-snowflake/tests/turu/test_snowflake_async.py +++ b/turu-snowflake/tests/turu/test_snowflake_async.py @@ -442,7 +442,7 @@ async def test_record_pandas_dataframe_with_limit_option( ), limit=2, ) as cursor: - expected = pd.DataFrame({"ID": list(range(1, 10))}, dtype="object") + expected = pd.DataFrame({"ID": [1, 2]}, dtype="object") assert_frame_equal(await cursor.fetch_pandas_all(), expected)