Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: write polars dataframe #113

Merged
merged 4 commits into from
Nov 12, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion influxdb_client_3/write_client/client/write_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
30 changes: 29 additions & 1 deletion tests/test_polars.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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)
Loading