fix: write polars dataframe (#113)

* fix: write polars dataframe * docs: update CHANGELOG.md * chore: add tests * docs: Update CHANGELOG.md --------- Co-authored-by: Jakub Bednář <jakub.bednar@gmail.com>
InfluxCommunity · Nov 12, 2024 · bc48b17 · bc48b17
1 parent 0303b3f
commit bc48b17
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 2 deletions.
diff --git 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
@@ -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
@@ -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)