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)