From bc48b1745393368417991a926b3c8d4c7aba1ae7 Mon Sep 17 00:00:00 2001 From: Tapanhaz <91151267+Tapanhaz@users.noreply.github.com> Date: Tue, 12 Nov 2024 09:37:31 +0530 Subject: [PATCH] fix: write polars dataframe (#113) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: write polars dataframe * docs: update CHANGELOG.md * chore: add tests * docs: Update CHANGELOG.md --------- Co-authored-by: Jakub Bednář --- CHANGELOG.md | 4 +++ .../write_client/client/write_api.py | 2 +- tests/test_polars.py | 30 ++++++++++++++++++- 3 files changed, 34 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c7a43a8..0c54ffc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,10 @@ ## 0.10.0 [unreleased] +### Bug Fixes + +1. [#113](https://github.com/InfluxCommunity/influxdb3-python/pull/113): Fix import error of `PolarsDataframeSerializer` in batching mode + ## 0.9.0 [2024-09-13] ### Features diff --git a/influxdb_client_3/write_client/client/write_api.py b/influxdb_client_3/write_client/client/write_api.py index 3519161..7181813 100644 --- a/influxdb_client_3/write_client/client/write_api.py +++ b/influxdb_client_3/write_client/client/write_api.py @@ -461,7 +461,7 @@ def _write_batching(self, bucket, org, data, precision, **kwargs) elif 'polars' in str(type(data)): - from influxdb_client_3.write_client.client.write.dataframe_serializer import PolarsDataframeSerializer + from influxdb_client_3.write_client.client.write.polars_dataframe_serializer import PolarsDataframeSerializer serializer = PolarsDataframeSerializer(data, self._point_settings, precision, self._write_options.batch_size, **kwargs) diff --git a/tests/test_polars.py b/tests/test_polars.py index dabde23..b1934bd 100644 --- a/tests/test_polars.py +++ b/tests/test_polars.py @@ -2,7 +2,7 @@ import unittest from unittest.mock import Mock -from influxdb_client_3 import PointSettings, InfluxDBClient3 +from influxdb_client_3 import PointSettings, InfluxDBClient3, write_client_options, WriteOptions from influxdb_client_3.write_client import WriteService from influxdb_client_3.write_client.client.write.polars_dataframe_serializer import polars_data_frame_to_list_of_points @@ -61,3 +61,31 @@ def test_write_polars(self): actual = self.client._write_api._write_service.post_write.call_args[1]['body'] self.assertEqual(b'measurement temperature=22.4 1722470400000000000\n' b'measurement temperature=21.8 1722474000000000000', actual) + + def test_write_polars_batching(self): + import polars as pl + df = pl.DataFrame({ + "time": pl.Series(["2024-08-01 00:00:00", "2024-08-01 01:00:00"]).str.to_datetime(time_unit='ns'), + "temperature": [22.4, 21.8], + }) + self.client = InfluxDBClient3( + host="localhost", + org="my_org", + database="my_db", + token="my_token", write_client_options=write_client_options( + write_options=WriteOptions(batch_size=2) + ) + ) + self.client._write_api._write_options = WriteOptions(batch_size=2) + self.client._write_api._write_service = Mock(spec=WriteService) + + self.client.write( + database="database", + record=df, + data_frame_measurement_name="measurement", + data_frame_timestamp_column="time", + ) + + actual = self.client._write_api._write_service.post_write.call_args[1]['body'] + self.assertEqual(b'measurement temperature=22.4 1722470400000000000\n' + b'measurement temperature=21.8 1722474000000000000', actual)