You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tsfile.apache.org by co...@apache.org on 2026/01/14 14:21:33 UTC
(tsfile) branch support_dataframe_to_tsfile created (now b2de41c4)
This is an automated email from the ASF dual-hosted git repository.
colinlee pushed a change to branch support_dataframe_to_tsfile
in repository https://gitbox.apache.org/repos/asf/tsfile.git
at b2de41c4 fix.
This branch includes the following new commits:
new ca1ad79e tmp code.
new dad7c61a support to_tsfile.
new b2de41c4 fix.
The 3 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(tsfile) 01/03: tmp code.
Posted by co...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
colinlee pushed a commit to branch support_dataframe_to_tsfile
in repository https://gitbox.apache.org/repos/asf/tsfile.git
commit ca1ad79efc8169b454c6f79adddff5645f51d83b
Author: ColinLee <sh...@163.com>
AuthorDate: Wed Jan 14 20:59:56 2026 +0800
tmp code.
---
python/tests/test_write_and_read.py | 225 ++++++++++++++++++++++++++++++++++-
python/tsfile/constants.py | 81 ++++++++++++-
python/tsfile/tsfile_cpp.pxd | 3 +-
python/tsfile/tsfile_py_cpp.pxd | 1 +
python/tsfile/tsfile_py_cpp.pyx | 182 ++++++++++++++++++++++++++--
python/tsfile/tsfile_table_writer.py | 72 ++++++++++-
python/tsfile/tsfile_writer.pyx | 30 +++--
python/tsfile/utils.py | 116 ++++++++++++++++++
8 files changed, 678 insertions(+), 32 deletions(-)
diff --git a/python/tests/test_write_and_read.py b/python/tests/test_write_and_read.py
index b327e2d3..ad8d698a 100644
--- a/python/tests/test_write_and_read.py
+++ b/python/tests/test_write_and_read.py
@@ -16,6 +16,7 @@
# under the License.
#
+import os
from datetime import date
import numpy as np
@@ -31,7 +32,7 @@ from tsfile import TimeseriesSchema
from tsfile import TsFileTableWriter
from tsfile import TsFileWriter, TsFileReader, ColumnCategory
from tsfile import to_dataframe
-from tsfile.exceptions import TableNotExistError, ColumnNotExistError, NotSupportedError
+from tsfile.exceptions import TableNotExistError, ColumnNotExistError, NotSupportedError, TypeMismatchError
def test_row_record_write_and_read():
@@ -985,7 +986,227 @@ def test_table_all_datatype_query_to_dataframe_variants():
os.remove(tsfile_path)
-import os
+def convert_to_nullable_types(df):
+ """
+ Convert DataFrame columns to nullable types to match returned DataFrame from to_dataframe.
+ This handles the fact that returned DataFrames use nullable types (Int64, Float64, etc.)
+ to support Null values.
+ """
+ df = df.copy()
+ for col in df.columns:
+ dtype = df[col].dtype
+ if dtype == 'int64':
+ df[col] = df[col].astype('Int64')
+ elif dtype == 'int32':
+ df[col] = df[col].astype('Int32')
+ elif dtype == 'float64':
+ df[col] = df[col].astype('Float64')
+ elif dtype == 'float32':
+ df[col] = df[col].astype('Float32')
+ elif dtype == 'bool':
+ df[col] = df[col].astype('boolean')
+ return df
+
+
+def test_write_dataframe_basic():
+ table = TableSchema("test_table",
+ [ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG),
+ ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD),
+ ColumnSchema("value2", TSDataType.INT64, ColumnCategory.FIELD)])
+ tsfile_path = "test_write_dataframe_basic.tsfile"
+ try:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+ with TsFileTableWriter(tsfile_path, table) as writer:
+ # Create DataFrame with 'time' column
+ df = pd.DataFrame({
+ 'time': [i for i in range(100)],
+ 'device': [f"device{i}" for i in range(100)],
+ 'value': [i * 1.5 for i in range(100)],
+ 'value2': [i * 10 for i in range(100)]
+ })
+ writer.write_dataframe(df)
+
+ df_read = to_dataframe(tsfile_path, table_name="test_table")
+ # Sort both DataFrames by time for comparison
+ df_read = df_read.sort_values('time').reset_index(drop=True)
+ df_sorted = convert_to_nullable_types(df.sort_values('time').reset_index(drop=True))
+ assert df_read.shape == (100, 4)
+ assert df_read["time"].equals(df_sorted["time"])
+ assert df_read["device"].equals(df_sorted["device"])
+ assert df_read["value"].equals(df_sorted["value"])
+ assert df_read["value2"].equals(df_sorted["value2"])
+ finally:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+
+def test_write_dataframe_with_index():
+ """Test write_dataframe using DataFrame index as time when no 'time' column exists."""
+ table = TableSchema("test_table",
+ [ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG),
+ ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD)])
+ tsfile_path = "test_write_dataframe_index.tsfile"
+ try:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+ with TsFileTableWriter(tsfile_path, table) as writer:
+ # Create DataFrame without 'time' column, use index
+ df = pd.DataFrame({
+ 'device': [f"device{i}" for i in range(50)],
+ 'value': [i * 2.5 for i in range(50)]
+ })
+ df.index = [i * 10 for i in range(50)] # Set index as timestamps
+ writer.write_dataframe(df)
+
+ # Verify by reading back
+ df_read = to_dataframe(tsfile_path, table_name="test_table")
+ # Sort both DataFrames by time for comparison
+ df_read = df_read.sort_values('time').reset_index(drop=True)
+ df_sorted = df.sort_index()
+ df_sorted = convert_to_nullable_types(df_sorted.reset_index(drop=True))
+ # Create time series from original index
+ time_series = pd.Series(df.sort_index().index.values, dtype='Int64')
+ assert df_read.shape == (50, 3)
+ # Compare time column with sorted index
+ assert df_read["time"].equals(time_series)
+ assert df_read["device"].equals(df_sorted["device"])
+ assert df_read["value"].equals(df_sorted["value"])
+ finally:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+
+def test_write_dataframe_case_insensitive():
+ """Test write_dataframe with case-insensitive column matching."""
+ table = TableSchema("test_table",
+ [ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG),
+ ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD)])
+ tsfile_path = "test_write_dataframe_case.tsfile"
+ try:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+ with TsFileTableWriter(tsfile_path, table) as writer:
+ # DataFrame with different case column names
+ df = pd.DataFrame({
+ 'Time': [i for i in range(30)], # Capital T
+ 'Device': [f"device{i}" for i in range(30)], # Capital D
+ 'VALUE': [i * 3.0 for i in range(30)] # All caps
+ })
+ writer.write_dataframe(df)
+
+ # Verify by reading back
+ df_read = to_dataframe(tsfile_path, table_name="test_table")
+ # Sort both DataFrames by time for comparison
+ df_read = df_read.sort_values('time').reset_index(drop=True)
+ df_sorted = convert_to_nullable_types(df.sort_values('Time').reset_index(drop=True))
+ assert df_read.shape == (30, 3)
+ assert df_read["time"].equals(df_sorted["Time"])
+ assert df_read["device"].equals(df_sorted["Device"])
+ assert df_read["value"].equals(df_sorted["VALUE"])
+ finally:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+
+def test_write_dataframe_column_not_in_schema():
+ """Test write_dataframe raises ColumnNotExistError when DataFrame has extra columns."""
+ table = TableSchema("test_table",
+ [ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG),
+ ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD)])
+ tsfile_path = "test_write_dataframe_extra_col.tsfile"
+ try:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+ with TsFileTableWriter(tsfile_path, table) as writer:
+ # DataFrame with extra column not in schema
+ df = pd.DataFrame({
+ 'time': [i for i in range(10)],
+ 'device': [f"device{i}" for i in range(10)],
+ 'value': [i * 1.0 for i in range(10)],
+ 'extra_column': [i for i in range(10)] # Not in schema
+ })
+ with pytest.raises(ColumnNotExistError) as exc_info:
+ writer.write_dataframe(df)
+ assert "extra_column" in str(exc_info.value)
+ finally:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+
+def test_write_dataframe_type_mismatch():
+ """Test write_dataframe raises TypeMismatchError when types are incompatible."""
+ table = TableSchema("test_table",
+ [ColumnSchema("value", TSDataType.STRING, ColumnCategory.FIELD)])
+ tsfile_path = "test_write_dataframe_type_mismatch.tsfile"
+ try:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+ with TsFileTableWriter(tsfile_path, table) as writer:
+ # DataFrame with incompatible type (INT64 vs STRING)
+ df = pd.DataFrame({
+ 'time': [i for i in range(10)],
+ 'value': [i for i in range(10)] # INT64, but schema expects STRING
+ })
+ with pytest.raises(TypeMismatchError) as exc_info:
+ writer.write_dataframe(df)
+ assert "Type mismatches" in str(exc_info.value)
+ finally:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+
+def test_write_dataframe_all_datatypes():
+ """Test write_dataframe with all supported data types."""
+ table = TableSchema("test_table",
+ [ColumnSchema("bool_col", TSDataType.BOOLEAN, ColumnCategory.FIELD),
+ ColumnSchema("int32_col", TSDataType.INT32, ColumnCategory.FIELD),
+ ColumnSchema("int64_col", TSDataType.INT64, ColumnCategory.FIELD),
+ ColumnSchema("float_col", TSDataType.FLOAT, ColumnCategory.FIELD),
+ ColumnSchema("double_col", TSDataType.DOUBLE, ColumnCategory.FIELD),
+ ColumnSchema("string_col", TSDataType.STRING, ColumnCategory.FIELD),
+ ColumnSchema("blob_col", TSDataType.BLOB, ColumnCategory.FIELD)])
+ tsfile_path = "test_write_dataframe_all_types.tsfile"
+ try:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+ with TsFileTableWriter(tsfile_path, table) as writer:
+ df = pd.DataFrame({
+ 'time': [i for i in range(50)],
+ 'bool_col': [i % 2 == 0 for i in range(50)],
+ 'int32_col': pd.Series([i for i in range(50)], dtype='int32'),
+ 'int64_col': [i * 10 for i in range(50)],
+ 'float_col': pd.Series([i * 1.5 for i in range(50)], dtype='float32'),
+ 'double_col': [i * 2.5 for i in range(50)],
+ 'string_col': [f"str{i}" for i in range(50)],
+ 'blob_col': [f"blob{i}".encode('utf-8') for i in range(50)]
+ })
+ writer.write_dataframe(df)
+
+ # Verify by reading back
+ df_read = to_dataframe(tsfile_path, table_name="test_table")
+ # Sort both DataFrames by time for comparison
+ df_read = df_read.sort_values('time').reset_index(drop=True)
+ df_sorted = convert_to_nullable_types(df.sort_values('time').reset_index(drop=True))
+ assert df_read.shape == (50, 8)
+ assert df_read["bool_col"].equals(df_sorted["bool_col"])
+ assert df_read["int32_col"].equals(df_sorted["int32_col"])
+ assert df_read["int64_col"].equals(df_sorted["int64_col"])
+ assert np.allclose(df_read["float_col"], df_sorted["float_col"])
+ assert np.allclose(df_read["double_col"], df_sorted["double_col"])
+ assert df_read["string_col"].equals(df_sorted["string_col"])
+ # BLOB comparison
+ for i in range(50):
+ assert df_read["blob_col"].iloc[i] == df_sorted["blob_col"].iloc[i]
+ finally:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
if __name__ == "__main__":
os.chdir(os.path.dirname(os.path.abspath(__file__)))
diff --git a/python/tsfile/constants.py b/python/tsfile/constants.py
index 7d1f5ff5..e6d9dd6b 100644
--- a/python/tsfile/constants.py
+++ b/python/tsfile/constants.py
@@ -16,6 +16,7 @@
# under the License.
#
from enum import unique, IntEnum
+import numpy as np
@unique
@@ -62,13 +63,13 @@ class TSDataType(IntEnum):
elif self == TSDataType.INT64:
return "Int64"
elif self == TSDataType.FLOAT:
- return "float32"
+ return "Float32"
elif self == TSDataType.DOUBLE:
- return "float64"
+ return "Float64"
elif self == TSDataType.TEXT or self == TSDataType.STRING:
return "object"
elif self == TSDataType.TIMESTAMP:
- return "int64"
+ return "Int64"
elif self == TSDataType.DATE:
return "object"
elif self == TSDataType.BLOB:
@@ -76,6 +77,80 @@ class TSDataType(IntEnum):
else:
raise ValueError(f"Unknown data type: {self}")
+ @classmethod
+ def from_pandas_datatype(cls, dtype):
+ if dtype is np.bool_:
+ return cls.BOOLEAN
+ elif dtype is np.int32:
+ return cls.INT32
+ elif dtype is np.int64:
+ return cls.INT64
+ elif dtype is np.float32:
+ return cls.FLOAT
+ elif dtype is np.float64:
+ return cls.DOUBLE
+ elif dtype is np.object_:
+ return cls.STRING
+
+ try:
+ import pandas as pd
+ if hasattr(pd, 'StringDtype') and isinstance(dtype, pd.StringDtype):
+ return cls.STRING
+ except (ImportError, AttributeError):
+ pass
+
+ if hasattr(dtype, 'type'):
+ dtype = dtype.type
+ if dtype is np.bool_:
+ return cls.BOOLEAN
+ elif dtype is np.int32:
+ return cls.INT32
+ elif dtype is np.int64:
+ return cls.INT64
+ elif dtype is np.float32:
+ return cls.FLOAT
+ elif dtype is np.float64:
+ return cls.DOUBLE
+ elif dtype is np.object_:
+ # object 类型默认返回 STRING，实际类型需要在 write_dataframe 中检查数据内容
+ return cls.STRING
+
+ dtype_str = str(dtype)
+
+ if 'stringdtype' in dtype_str.lower() or dtype_str.startswith('string'):
+ return cls.STRING
+
+ dtype_map = {
+ 'bool': cls.BOOLEAN,
+ 'boolean': cls.BOOLEAN,
+ 'int32': cls.INT32,
+ 'Int32': cls.INT32,
+ 'int64': cls.INT64,
+ 'Int64': cls.INT64,
+ 'float32': cls.FLOAT,
+ 'float64': cls.DOUBLE,
+ 'bytes': cls.BLOB,
+ 'object': cls.STRING,
+ 'string': cls.STRING,
+ }
+
+ if dtype_str in dtype_map:
+ return dtype_map[dtype_str]
+
+ dtype_lower = dtype_str.lower()
+ if dtype_lower in dtype_map:
+ return dtype_map[dtype_lower]
+
+ if 'object_' in dtype_lower or dtype_str == "<class 'numpy.object_'>":
+ return cls.STRING
+
+ if dtype_str.startswith('datetime64'):
+ return cls.TIMESTAMP
+
+ return cls.STRING
+
+
+
@unique
class TSEncoding(IntEnum):
diff --git a/python/tsfile/tsfile_cpp.pxd b/python/tsfile/tsfile_cpp.pxd
index 40bff4eb..ab915fef 100644
--- a/python/tsfile/tsfile_cpp.pxd
+++ b/python/tsfile/tsfile_cpp.pxd
@@ -137,7 +137,8 @@ cdef extern from "./tsfile_cwrapper.h":
TSDataType * data_types,
int column_num, int max_rows);
- Tablet tablet_new(const char** column_names, TSDataType * data_types, int column_num);
+ Tablet tablet_new(char** column_name_list, TSDataType* data_types,
+ uint32_t column_num, uint32_t max_rows);
ErrorCode tablet_add_timestamp(Tablet tablet, uint32_t row_index, int64_t timestamp);
ErrorCode tablet_add_value_by_index_int64_t(Tablet tablet, uint32_t row_index, uint32_t column_index,
diff --git a/python/tsfile/tsfile_py_cpp.pxd b/python/tsfile/tsfile_py_cpp.pxd
index e44bb588..9ce2f90d 100644
--- a/python/tsfile/tsfile_py_cpp.pxd
+++ b/python/tsfile/tsfile_py_cpp.pxd
@@ -33,6 +33,7 @@ cdef public api DeviceSchema* to_c_device_schema(object py_schema)
cdef public api ColumnSchema* to_c_column_schema(object py_schema)
cdef public api TableSchema* to_c_table_schema(object py_schema)
cdef public api Tablet to_c_tablet(object tablet)
+cdef public api Tablet dataframe_to_c_tablet(object target_name, object dataframe)
cdef public api TsRecord to_c_record(object row_record)
cdef public api void free_c_table_schema(TableSchema* c_schema)
cdef public api void free_c_column_schema(ColumnSchema* c_schema)
diff --git a/python/tsfile/tsfile_py_cpp.pyx b/python/tsfile/tsfile_py_cpp.pyx
index d9924d7a..84c550f9 100644
--- a/python/tsfile/tsfile_py_cpp.pyx
+++ b/python/tsfile/tsfile_py_cpp.pyx
@@ -19,6 +19,9 @@
from .date_utils import parse_date_to_int
from .tsfile_cpp cimport *
+import pandas as pd
+import numpy as np
+
from libc.stdlib cimport free
from libc.stdlib cimport malloc
from libc.string cimport strdup
@@ -220,7 +223,7 @@ cdef Tablet to_c_tablet(object tablet):
cdef TSDataType * column_types
cdef bytes row_bytes
cdef char *raw_str
- cdef const char* str_ptr
+ cdef const char * str_ptr
cdef Py_ssize_t raw_len
if tablet.get_target_name() is not None:
@@ -293,7 +296,7 @@ cdef Tablet to_c_tablet(object tablet):
for row in range(max_row_num):
if value[row] is not None:
py_value = value[row]
- str_ptr = PyUnicode_AsUTF8AndSize(py_value, &raw_len)
+ str_ptr = PyUnicode_AsUTF8AndSize(py_value, &raw_len)
tablet_add_value_by_index_string_with_len(ctablet, row, col, str_ptr, raw_len)
elif data_type == TS_DATATYPE_BLOB:
@@ -304,13 +307,166 @@ cdef Tablet to_c_tablet(object tablet):
return ctablet
+cdef TSDataType pandas_dtype_to_ts_data_type(object dtype):
+ return to_c_data_type(TSDataTypePy.from_pandas_datatype(dtype))
+
+cdef Tablet dataframe_to_c_tablet(object target_name, object dataframe):
+ cdef Tablet ctablet
+ cdef int max_row_num
+ cdef TSDataType data_type
+ cdef int64_t timestamp
+ cdef const char * device_id_c = NULL
+ cdef char** columns_names
+ cdef TSDataType * columns_types
+ cdef char *raw_str
+ cdef const char * str_ptr
+ cdef Py_ssize_t raw_len
+ cdef int column_num
+ cdef int i, row
+ cdef object value
+ cdef object py_value
+ cdef object value_bytes
+
+ device_id_bytes = PyUnicode_AsUTF8String(target_name.lower())
+ device_id_c = device_id_bytes
+ df_columns = list(dataframe.columns)
+ use_id_as_time = False
+ time_column_name = None
+
+ # Find column with lowercase name 'time'
+
+
+ for col in df_columns:
+ if col.lower() == 'time':
+ time_column_name = col
+ break
+
+ if time_column_name is None:
+ use_id_as_time = True
+
+ data_columns = [col for col in df_columns if col.lower() != 'time']
+ column_num = len(data_columns)
+
+ if column_num == 0:
+ raise ValueError("DataFrame must have at least one data column besides 'time'")
+
+ max_row_num = len(dataframe)
+
+ column_types_list = []
+ for col_name in data_columns:
+ pandas_dtype = dataframe[col_name].dtype
+ ds_type = pandas_dtype_to_ts_data_type(pandas_dtype)
+ if ds_type == TS_DATATYPE_STRING:
+ column_series = dataframe[col_name]
+ first_valid_idx = column_series.first_valid_index()
+ if first_valid_idx is not None:
+ first_value = column_series[first_valid_idx]
+ if isinstance(first_value, bytes):
+ ds_type = TS_DATATYPE_BLOB
+ column_types_list.append(ds_type)
+
+ columns_names = <char**> malloc(sizeof(char *) * column_num)
+ columns_types = <TSDataType *> malloc(sizeof(TSDataType) * column_num)
+
+ for i in range(column_num):
+ columns_names[i] = strdup(data_columns[i].lower().encode('utf-8'))
+ columns_types[i] = column_types_list[i]
+
+ ctablet = _tablet_new_with_target_name(device_id_c, columns_names, columns_types, column_num,
+ max_row_num)
+
+ free(columns_types)
+ for i in range(column_num):
+ free(columns_names[i])
+ free(columns_names)
+
+ if use_id_as_time:
+ for row in range(max_row_num):
+ timestamp_py = dataframe.index[row]
+ if pd.isna(timestamp_py):
+ continue
+ timestamp = <int64_t> timestamp_py
+ tablet_add_timestamp(ctablet, row, timestamp)
+ else:
+ time_values = dataframe[time_column_name].values
+ for row in range(max_row_num):
+ timestamp_py = time_values[row]
+ if pd.isna(timestamp_py):
+ continue
+ timestamp = <int64_t> timestamp_py
+ tablet_add_timestamp(ctablet, row, timestamp)
+
+ for col in range(column_num):
+ col_name = data_columns[col]
+ data_type = column_types_list[col]
+ column_values = dataframe[col_name].values
+
+ # BOOLEAN
+ if data_type == TS_DATATYPE_BOOLEAN:
+ for row in range(max_row_num):
+ value = column_values[row]
+ if not pd.isna(value):
+ tablet_add_value_by_index_bool(ctablet, row, col, <bint> value)
+ # INT32
+ elif data_type == TS_DATATYPE_INT32:
+ for row in range(max_row_num):
+ value = column_values[row]
+ if not pd.isna(value):
+ tablet_add_value_by_index_int32_t(ctablet, row, col, <int32_t> value)
+ # INT64
+ elif data_type == TS_DATATYPE_INT64 or data_type == TS_DATATYPE_TIMESTAMP:
+ for row in range(max_row_num):
+ value = column_values[row]
+ if not pd.isna(value):
+ tablet_add_value_by_index_int64_t(ctablet, row, col, <int64_t> value)
+ # FLOAT
+ elif data_type == TS_DATATYPE_FLOAT:
+ for row in range(max_row_num):
+ value = column_values[row]
+ if not pd.isna(value):
+ tablet_add_value_by_index_float(ctablet, row, col, <float> value)
+ # DOUBLE
+ elif data_type == TS_DATATYPE_DOUBLE:
+ for row in range(max_row_num):
+ value = column_values[row]
+ if not pd.isna(value):
+ tablet_add_value_by_index_double(ctablet, row, col, <double> value)
+ # DATE
+ elif data_type == TS_DATATYPE_DATE:
+ for row in range(max_row_num):
+ value = column_values[row]
+ if not pd.isna(value):
+ tablet_add_value_by_index_int32_t(ctablet, row, col, parse_date_to_int(value))
+ # STRING or TEXT
+ elif data_type == TS_DATATYPE_STRING or data_type == TS_DATATYPE_TEXT:
+ for row in range(max_row_num):
+ value = column_values[row]
+ if not pd.isna(value):
+ py_value = str(value)
+ str_ptr = PyUnicode_AsUTF8AndSize(py_value, &raw_len)
+ tablet_add_value_by_index_string_with_len(ctablet, row, col, str_ptr, raw_len)
+ # BLOB
+ elif data_type == TS_DATATYPE_BLOB:
+ for row in range(max_row_num):
+ value = column_values[row]
+ if not pd.isna(value):
+ if isinstance(value, bytes):
+ PyBytes_AsStringAndSize(value, &raw_str, &raw_len)
+ tablet_add_value_by_index_string_with_len(ctablet, row, col, raw_str, raw_len)
+ else:
+ value_bytes = bytes(value)
+ PyBytes_AsStringAndSize(value_bytes, &raw_str, &raw_len)
+ tablet_add_value_by_index_string_with_len(ctablet, row, col, raw_str, raw_len)
+
+ return ctablet
+
cdef TsRecord to_c_record(object row_record):
cdef int field_num = row_record.get_fields_num()
cdef int64_t timestamp = <int64_t> row_record.get_timestamp()
cdef bytes device_id_bytes = PyUnicode_AsUTF8String(row_record.get_device_id())
- cdef const char* device_id = device_id_bytes
- cdef const char* str_ptr
- cdef char* blob_ptr
+ cdef const char * device_id = device_id_bytes
+ cdef const char * str_ptr
+ cdef char * blob_ptr
cdef Py_ssize_t str_len
cdef TsRecord record
cdef int i
@@ -320,9 +476,11 @@ cdef TsRecord to_c_record(object row_record):
field = row_record.get_fields()[i]
data_type = to_c_data_type(field.get_data_type())
if data_type == TS_DATATYPE_BOOLEAN:
- _insert_data_into_ts_record_by_name_bool(record, PyUnicode_AsUTF8(field.get_field_name()), field.get_bool_value())
+ _insert_data_into_ts_record_by_name_bool(record, PyUnicode_AsUTF8(field.get_field_name()),
+ field.get_bool_value())
elif data_type == TS_DATATYPE_INT32 or data_type == TS_DATATYPE_DATE:
- _insert_data_into_ts_record_by_name_int32_t(record, PyUnicode_AsUTF8(field.get_field_name()), field.get_int_value())
+ _insert_data_into_ts_record_by_name_int32_t(record, PyUnicode_AsUTF8(field.get_field_name()),
+ field.get_int_value())
elif data_type == TS_DATATYPE_INT64:
_insert_data_into_ts_record_by_name_int64_t(record, PyUnicode_AsUTF8(field.get_field_name()),
field.get_long_value())
@@ -333,15 +491,17 @@ cdef TsRecord to_c_record(object row_record):
_insert_data_into_ts_record_by_name_double(record, PyUnicode_AsUTF8(field.get_field_name()),
field.get_double_value())
elif data_type == TS_DATATYPE_FLOAT:
- _insert_data_into_ts_record_by_name_float(record, PyUnicode_AsUTF8(field.get_field_name()), field.get_float_value())
+ _insert_data_into_ts_record_by_name_float(record, PyUnicode_AsUTF8(field.get_field_name()),
+ field.get_float_value())
elif data_type == TS_DATATYPE_TEXT or data_type == TS_DATATYPE_STRING:
- str_ptr = PyUnicode_AsUTF8AndSize(field.get_string_value(), &str_len)
- _insert_data_into_ts_record_by_name_string_with_len(record, PyUnicode_AsUTF8(field.get_field_name()), str_ptr, str_len)
+ str_ptr = PyUnicode_AsUTF8AndSize(field.get_string_value(), &str_len)
+ _insert_data_into_ts_record_by_name_string_with_len(record, PyUnicode_AsUTF8(field.get_field_name()),
+ str_ptr, str_len)
elif data_type == TS_DATATYPE_BLOB:
if PyBytes_AsStringAndSize(field.get_string_value(), &blob_ptr, &str_len) < 0:
raise ValueError("blob not legal")
_insert_data_into_ts_record_by_name_string_with_len(record, PyUnicode_AsUTF8(field.get_field_name()),
- <const char*> blob_ptr, <uint32_t>str_len)
+ <const char *> blob_ptr, <uint32_t> str_len)
return record
# Free c structs' space
diff --git a/python/tsfile/tsfile_table_writer.py b/python/tsfile/tsfile_table_writer.py
index 28193360..54f759c9 100644
--- a/python/tsfile/tsfile_table_writer.py
+++ b/python/tsfile/tsfile_table_writer.py
@@ -15,9 +15,12 @@
# specific language governing permissions and limitations
# under the License.
#
+import pandas as pd
from tsfile import TableSchema, Tablet, TableNotExistError
from tsfile import TsFileWriter
+from tsfile.constants import TSDataType
+from tsfile.exceptions import ColumnNotExistError, TypeMismatchError
class TsFileTableWriter:
@@ -31,7 +34,7 @@ class TsFileTableWriter:
according to that schema, and serialize this data into a TsFile.
"""
- def __init__(self, path: str, table_schema: TableSchema, memory_threshold = 128 * 1024 * 1024):
+ def __init__(self, path: str, table_schema: TableSchema, memory_threshold=128 * 1024 * 1024):
"""
:param path: The path of tsfile, will create if it doesn't exist.
:param table_schema: describes the schema of the tables they want to write.
@@ -39,7 +42,7 @@ class TsFileTableWriter:
"""
self.writer = TsFileWriter(path, memory_threshold)
self.writer.register_table(table_schema)
- self.exclusive_table_name_ = table_schema.get_table_name()
+ self.tableSchema = table_schema
def write_table(self, tablet: Tablet):
"""
@@ -49,11 +52,72 @@ class TsFileTableWriter:
:raise: TableNotExistError if table does not exist or tablet's table_name does not match tableschema.
"""
if tablet.get_target_name() is None:
- tablet.set_table_name(self.exclusive_table_name_)
- elif self.exclusive_table_name_ is not None and tablet.get_target_name() != self.exclusive_table_name_:
+ tablet.set_table_name(self.tableSchema.get_table_name())
+ elif (self.tableSchema.get_table_name() is not None
+ and tablet.get_target_name() != self.tableSchema.get_table_name()):
raise TableNotExistError
self.writer.write_table(tablet)
+ def write_dataframe(self, dataframe: pd.DataFrame):
+ """
+ Write a pandas DataFrame into table in tsfile.
+ :param dataframe: pandas DataFrame with 'time' column and data columns matching schema.
+ :return: no return value.
+ :raise: ColumnNotExistError if DataFrame columns don't match schema.
+ :raise: TypeMismatchError if DataFrame column types are incompatible with schema.
+ """
+ if dataframe is None or dataframe.empty:
+ pass
+
+ # Create mapping from lowercase column name to original column name
+ df_column_name_map = {col.lower(): col for col in dataframe.columns if col.lower() != 'time'}
+ df_columns = list(df_column_name_map.keys())
+
+ schema_column_names = set(self.tableSchema.get_column_names())
+ df_columns_set = set(df_columns)
+
+ extra_columns = df_columns_set - schema_column_names
+ if extra_columns:
+ raise ColumnNotExistError(
+ code=50,
+ context=f"DataFrame has columns not in schema: {', '.join(sorted(extra_columns))}"
+ )
+
+ schema_column_map = {
+ col.get_column_name(): col for col in self.tableSchema.get_columns()
+ }
+
+ type_mismatches = []
+ for col_name in df_columns:
+ df_col_name_original = df_column_name_map[col_name]
+
+ df_dtype = dataframe[df_col_name_original].dtype
+ df_ts_type = TSDataType.from_pandas_datatype(df_dtype)
+
+ if df_ts_type == TSDataType.STRING and (df_dtype == 'object' or str(df_dtype) == "<class 'numpy.object_'>"):
+ column_series = dataframe[df_col_name_original]
+ first_valid_idx = column_series.first_valid_index()
+ if first_valid_idx is not None:
+ first_value = column_series[first_valid_idx]
+ if isinstance(first_value, bytes):
+ df_ts_type = TSDataType.BLOB
+
+ schema_col = schema_column_map[col_name]
+ expected_ts_type = schema_col.get_data_type()
+
+ if df_ts_type != expected_ts_type:
+ type_mismatches.append(
+ f"Column '{col_name}': expected {expected_ts_type.name}, got {df_ts_type.name}"
+ )
+
+ if type_mismatches:
+ raise TypeMismatchError(
+ code=27,
+ context=f"Type mismatches: {'; '.join(type_mismatches)}"
+ )
+
+ self.writer.write_dataframe(self.tableSchema.get_table_name(), dataframe)
+
def close(self):
"""
Close TsFileTableWriter and will flush data automatically.
diff --git a/python/tsfile/tsfile_writer.pyx b/python/tsfile/tsfile_writer.pyx
index 20199195..c558984e 100644
--- a/python/tsfile/tsfile_writer.pyx
+++ b/python/tsfile/tsfile_writer.pyx
@@ -15,21 +15,21 @@
# specific language governing permissions and limitations
# under the License.
#
-
-#cython: language_level=3
-
-from .tsfile_cpp cimport *
-from .tsfile_py_cpp cimport *
+import pandas
from tsfile.row_record import RowRecord
-from tsfile.schema import TimeseriesSchema as TimeseriesSchemaPy, DeviceSchema as DeviceSchemaPy
from tsfile.schema import TableSchema as TableSchemaPy
+from tsfile.schema import TimeseriesSchema as TimeseriesSchemaPy, DeviceSchema as DeviceSchemaPy
from tsfile.tablet import Tablet as TabletPy
+from .tsfile_cpp cimport *
+from .tsfile_py_cpp cimport *
+
+#cython: language_level=3
cdef class TsFileWriterPy:
cdef TsFileWriter writer
- def __init__(self, pathname:str, memory_threshold:int = 128 * 1024 * 1024):
+ def __init__(self, pathname: str, memory_threshold: int = 128 * 1024 * 1024):
self.writer = tsfile_writer_new_c(pathname, memory_threshold)
def register_timeseries(self, device_name : str, timeseries_schema : TimeseriesSchemaPy):
@@ -38,7 +38,7 @@ cdef class TsFileWriterPy:
device_name: device name of the timeseries
timeseries_schema: measurement's name/datatype/encoding/compressor
"""
- cdef TimeseriesSchema* c_schema = to_c_timeseries_schema(timeseries_schema)
+ cdef TimeseriesSchema * c_schema = to_c_timeseries_schema(timeseries_schema)
cdef ErrorCode errno
try:
errno = tsfile_writer_register_timeseries_py_cpp(self.writer, device_name, c_schema)
@@ -51,7 +51,7 @@ cdef class TsFileWriterPy:
Register a device with tsfile writer.
device_schema: the device definition, including device_name, some measurements' schema.
"""
- cdef DeviceSchema* device_schema_c = to_c_device_schema(device_schema)
+ cdef DeviceSchema * device_schema_c = to_c_device_schema(device_schema)
cdef ErrorCode errno
try:
errno = tsfile_writer_register_device_py_cpp(self.writer, device_schema_c)
@@ -64,7 +64,7 @@ cdef class TsFileWriterPy:
Register a table with tsfile writer.
table_schema: the table definition, include table_name, columns' schema.
"""
- cdef TableSchema* c_schema = to_c_table_schema(table_schema)
+ cdef TableSchema * c_schema = to_c_table_schema(table_schema)
cdef ErrorCode errno
try:
errno = tsfile_writer_register_table_py_cpp(self.writer, c_schema)
@@ -86,6 +86,15 @@ cdef class TsFileWriterPy:
finally:
free_c_tablet(ctablet)
+ def write_dataframe(self, target_table: str, dataframe: pandas.DataFrame):
+ cdef Tablet ctablet = dataframe_to_c_tablet(target_table, dataframe)
+ cdef ErrorCode errno
+ try:
+ errno = _tsfile_writer_write_table(self.writer, ctablet)
+ check_error(errno)
+ finally:
+ free_c_tablet(ctablet)
+
def write_row_record(self, record : RowRecord):
"""
Write a record into tsfile with tsfile writer.
@@ -143,4 +152,3 @@ cdef class TsFileWriterPy:
def __exit__(self, exc_type, exc_val, exc_tb):
self.close()
-
diff --git a/python/tsfile/utils.py b/python/tsfile/utils.py
index d27a0fae..a8ed5601 100644
--- a/python/tsfile/utils.py
+++ b/python/tsfile/utils.py
@@ -20,9 +20,11 @@ from typing import Optional
import numpy as np
import pandas as pd
+from pandas.core.dtypes.common import is_integer_dtype
from tsfile.exceptions import TableNotExistError, ColumnNotExistError
from tsfile.tsfile_reader import TsFileReaderPy
+from tsfile import ColumnSchema, TableSchema, ColumnCategory, TSDataType, TsFileTableWriter
def to_dataframe(file_path: str,
@@ -159,3 +161,117 @@ def to_dataframe(file_path: str,
return df
else:
return pd.DataFrame()
+
+
+def to_tsfile(dataframe: pd.DataFrame,
+ file_path: str,
+ table_name: Optional[str] = None,
+ time_column: Optional[str] = None,
+ tag_column: Optional[list[str]] = None,
+ ):
+ """
+ Write a pandas DataFrame to a TsFile by inferring the table schema from the DataFrame.
+
+ This function automatically infers the table schema based on the DataFrame's column
+ names and data types, then writes the data to a TsFile.
+
+ Parameters
+ ----------
+ dataframe : pd.DataFrame
+ The pandas DataFrame to write to TsFile.
+ - If a 'time' column (case-insensitive) exists, it will be used as the time column.
+ - Otherwise, the DataFrame index will be used as timestamps.
+ - All other columns will be treated as data columns.
+
+ file_path : str
+ Path to the TsFile to write. Will be created if it doesn't exist.
+
+ table_name : Optional[str], default None
+ Name of the table. If None, defaults to "table".
+
+ time_column : Optional[str], default None
+ Name of the time column. If None, will look for a column named 'time' (case-insensitive),
+ or use the DataFrame index if no 'time' column is found.
+
+ tag_column : Optional[list[str]], default None
+ List of column names to be treated as TAG columns. All other columns will be FIELD columns.
+ If None, all columns are treated as FIELD columns.
+
+ Returns
+ -------
+ None
+
+ Raises
+ ------
+ ValueError
+ If the DataFrame is empty or has no data columns.
+ """
+ if dataframe is None or dataframe.empty:
+ raise ValueError("DataFrame cannot be None or empty")
+
+ # Determine table name
+ if table_name is None:
+ table_name = "table"
+
+ # Determine time column
+ time_col_name = None
+ if time_column is not None:
+ # Check if specified time column exists
+ if time_column not in dataframe.columns:
+ raise ValueError(f"Time column '{time_column}' not found in DataFrame")
+ # Check if time column is integer type (int64 or int)
+ if not is_integer_dtype(dataframe[time_column].dtype):
+ raise ValueError(f"Time column '{time_column}' must be integer type (int64 or int), got {dataframe[time_column].dtype}")
+ time_col_name = time_column
+ else:
+ # Look for 'time' column (case-insensitive)
+ for col in dataframe.columns:
+ if col.lower() == 'time':
+ # Check if time column is integer type
+ if is_integer_dtype(dataframe[col].dtype):
+ time_col_name = col
+ break
+
+ # Get data columns (excluding time column)
+ data_columns = [col for col in dataframe.columns if col != time_col_name]
+
+ if len(data_columns) == 0:
+ raise ValueError("DataFrame must have at least one data column besides the time column")
+
+ # Normalize tag_column list (convert to lowercase for comparison)
+ tag_columns_lower = []
+ if tag_column is not None:
+ for tag_col in tag_column:
+ if tag_col not in dataframe.columns:
+ raise ValueError(f"Tag column '{tag_col}' not found in DataFrame")
+ tag_columns_lower.append(tag_col.lower())
+
+ # Infer schema from DataFrame columns
+ column_schemas = []
+ for col_name in data_columns:
+ # Infer data type from DataFrame column
+ col_dtype = dataframe[col_name].dtype
+ ts_data_type = TSDataType.from_pandas_datatype(col_dtype)
+
+ # Determine category (TAG or FIELD)
+ if col_name.lower() in tag_columns_lower:
+ category = ColumnCategory.TAG
+ else:
+ category = ColumnCategory.FIELD
+
+ column_schemas.append(ColumnSchema(col_name, ts_data_type, category))
+
+ # Create table schema
+ table_schema = TableSchema(table_name, column_schemas)
+
+ # Prepare DataFrame for writing
+ # If time column is specified or found, rename it to 'time' if needed
+ # If no time column is found, write_dataframe will handle it by using index
+ df_to_write = dataframe.copy()
+ if time_col_name is not None and time_col_name != 'time':
+ # Rename time column to 'time' for consistency
+ df_to_write = df_to_write.rename(columns={time_col_name: 'time'})
+
+ # Write DataFrame to TsFile
+ with TsFileTableWriter(file_path, table_schema) as writer:
+ writer.write_dataframe(df_to_write)
\ No newline at end of file
(tsfile) 02/03: support to_tsfile.
Posted by co...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
colinlee pushed a commit to branch support_dataframe_to_tsfile
in repository https://gitbox.apache.org/repos/asf/tsfile.git
commit dad7c61a3d7a8144ec66b72283c6312c7c418e61
Author: ColinLee <sh...@163.com>
AuthorDate: Wed Jan 14 21:35:53 2026 +0800
support to_tsfile.
---
python/tests/test_dataframe.py | 276 +++++++++++++++++++++
python/tests/test_to_tsfile.py | 376 +++++++++++++++++++++++++++++
python/tests/test_write_and_read.py | 455 +----------------------------------
python/tsfile/tsfile_table_writer.py | 3 +-
python/tsfile/utils.py | 13 +-
5 files changed, 670 insertions(+), 453 deletions(-)
diff --git a/python/tests/test_dataframe.py b/python/tests/test_dataframe.py
new file mode 100644
index 00000000..9c57e154
--- /dev/null
+++ b/python/tests/test_dataframe.py
@@ -0,0 +1,276 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+import os
+
+import numpy as np
+import pandas as pd
+import pytest
+from pandas.core.dtypes.common import is_integer_dtype
+
+from tsfile import ColumnSchema, TableSchema, TSDataType
+from tsfile import TsFileTableWriter, ColumnCategory
+from tsfile import to_dataframe
+from tsfile.exceptions import ColumnNotExistError, TypeMismatchError
+
+
+def convert_to_nullable_types(df):
+ """
+ Convert DataFrame columns to nullable types to match returned DataFrame from to_dataframe.
+ This handles the fact that returned DataFrames use nullable types (Int64, Float64, etc.)
+ to support Null values.
+ """
+ df = df.copy()
+ for col in df.columns:
+ dtype = df[col].dtype
+ if dtype == 'int64':
+ df[col] = df[col].astype('Int64')
+ elif dtype == 'int32':
+ df[col] = df[col].astype('Int32')
+ elif dtype == 'float64':
+ df[col] = df[col].astype('Float64')
+ elif dtype == 'float32':
+ df[col] = df[col].astype('Float32')
+ elif dtype == 'bool':
+ df[col] = df[col].astype('boolean')
+ return df
+
+
+def test_write_dataframe_basic():
+ """Test basic write_dataframe functionality."""
+ table = TableSchema("test_table",
+ [ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG),
+ ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD),
+ ColumnSchema("value2", TSDataType.INT64, ColumnCategory.FIELD)])
+ tsfile_path = "test_write_dataframe_basic.tsfile"
+ try:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+ with TsFileTableWriter(tsfile_path, table) as writer:
+ # Create DataFrame with 'time' column
+ df = pd.DataFrame({
+ 'time': [i for i in range(100)],
+ 'device': [f"device{i}" for i in range(100)],
+ 'value': [i * 1.5 for i in range(100)],
+ 'value2': [i * 10 for i in range(100)]
+ })
+ writer.write_dataframe(df)
+
+ df_read = to_dataframe(tsfile_path, table_name="test_table")
+ # Sort both DataFrames by time for comparison
+ df_read = df_read.sort_values('time').reset_index(drop=True)
+ df_sorted = convert_to_nullable_types(df.sort_values('time').reset_index(drop=True))
+ assert df_read.shape == (100, 4)
+ assert df_read["time"].equals(df_sorted["time"])
+ assert df_read["device"].equals(df_sorted["device"])
+ assert df_read["value"].equals(df_sorted["value"])
+ assert df_read["value2"].equals(df_sorted["value2"])
+ finally:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+
+def test_write_dataframe_with_index():
+ """Test write_dataframe using DataFrame index as time when no 'time' column exists."""
+ table = TableSchema("test_table",
+ [ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG),
+ ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD)])
+ tsfile_path = "test_write_dataframe_index.tsfile"
+ try:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+ with TsFileTableWriter(tsfile_path, table) as writer:
+ # Create DataFrame without 'time' column, use index
+ df = pd.DataFrame({
+ 'device': [f"device{i}" for i in range(50)],
+ 'value': [i * 2.5 for i in range(50)]
+ })
+ df.index = [i * 10 for i in range(50)] # Set index as timestamps
+ writer.write_dataframe(df)
+
+ # Verify by reading back
+ df_read = to_dataframe(tsfile_path, table_name="test_table")
+ # Sort both DataFrames by time for comparison
+ df_read = df_read.sort_values('time').reset_index(drop=True)
+ df_sorted = df.sort_index()
+ df_sorted = convert_to_nullable_types(df_sorted.reset_index(drop=True))
+ # Create time series from original index
+ time_series = pd.Series(df.sort_index().index.values, dtype='Int64')
+ assert df_read.shape == (50, 3)
+ # Compare time column with sorted index
+ assert df_read["time"].equals(time_series)
+ assert df_read["device"].equals(df_sorted["device"])
+ assert df_read["value"].equals(df_sorted["value"])
+ finally:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+
+def test_write_dataframe_case_insensitive():
+ """Test write_dataframe with case-insensitive column matching."""
+ table = TableSchema("test_table",
+ [ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG),
+ ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD)])
+ tsfile_path = "test_write_dataframe_case.tsfile"
+ try:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+ with TsFileTableWriter(tsfile_path, table) as writer:
+ # DataFrame with different case column names
+ df = pd.DataFrame({
+ 'Time': [i for i in range(30)], # Capital T
+ 'Device': [f"device{i}" for i in range(30)], # Capital D
+ 'VALUE': [i * 3.0 for i in range(30)] # All caps
+ })
+ writer.write_dataframe(df)
+
+ # Verify by reading back
+ df_read = to_dataframe(tsfile_path, table_name="test_table")
+ # Sort both DataFrames by time for comparison
+ df_read = df_read.sort_values('time').reset_index(drop=True)
+ df_sorted = convert_to_nullable_types(df.sort_values('Time').reset_index(drop=True))
+ assert df_read.shape == (30, 3)
+ assert df_read["time"].equals(df_sorted["Time"])
+ assert df_read["device"].equals(df_sorted["Device"])
+ assert df_read["value"].equals(df_sorted["VALUE"])
+ finally:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+
+def test_write_dataframe_column_not_in_schema():
+ """Test write_dataframe raises ColumnNotExistError when DataFrame has extra columns."""
+ table = TableSchema("test_table",
+ [ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG),
+ ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD)])
+ tsfile_path = "test_write_dataframe_extra_col.tsfile"
+ try:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+ with TsFileTableWriter(tsfile_path, table) as writer:
+ # DataFrame with extra column not in schema
+ df = pd.DataFrame({
+ 'time': [i for i in range(10)],
+ 'device': [f"device{i}" for i in range(10)],
+ 'value': [i * 1.0 for i in range(10)],
+ 'extra_column': [i for i in range(10)] # Not in schema
+ })
+ with pytest.raises(ColumnNotExistError) as exc_info:
+ writer.write_dataframe(df)
+ assert "extra_column" in str(exc_info.value)
+ finally:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+
+def test_write_dataframe_type_mismatch():
+ """Test write_dataframe raises TypeMismatchError when types are incompatible."""
+ table = TableSchema("test_table",
+ [ColumnSchema("value", TSDataType.STRING, ColumnCategory.FIELD)])
+ tsfile_path = "test_write_dataframe_type_mismatch.tsfile"
+ try:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+ with TsFileTableWriter(tsfile_path, table) as writer:
+ # DataFrame with incompatible type (INT64 vs STRING)
+ df = pd.DataFrame({
+ 'time': [i for i in range(10)],
+ 'value': [i for i in range(10)] # INT64, but schema expects STRING
+ })
+ with pytest.raises(TypeMismatchError) as exc_info:
+ writer.write_dataframe(df)
+ assert "Type mismatches" in str(exc_info.value)
+ finally:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+
+def test_write_dataframe_all_datatypes():
+ """Test write_dataframe with all supported data types."""
+ table = TableSchema("test_table",
+ [ColumnSchema("bool_col", TSDataType.BOOLEAN, ColumnCategory.FIELD),
+ ColumnSchema("int32_col", TSDataType.INT32, ColumnCategory.FIELD),
+ ColumnSchema("int64_col", TSDataType.INT64, ColumnCategory.FIELD),
+ ColumnSchema("float_col", TSDataType.FLOAT, ColumnCategory.FIELD),
+ ColumnSchema("double_col", TSDataType.DOUBLE, ColumnCategory.FIELD),
+ ColumnSchema("string_col", TSDataType.STRING, ColumnCategory.FIELD),
+ ColumnSchema("blob_col", TSDataType.BLOB, ColumnCategory.FIELD)])
+ tsfile_path = "test_write_dataframe_all_types.tsfile"
+ try:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+ with TsFileTableWriter(tsfile_path, table) as writer:
+ df = pd.DataFrame({
+ 'time': [i for i in range(50)],
+ 'bool_col': [i % 2 == 0 for i in range(50)],
+ 'int32_col': pd.Series([i for i in range(50)], dtype='int32'),
+ 'int64_col': [i * 10 for i in range(50)],
+ 'float_col': pd.Series([i * 1.5 for i in range(50)], dtype='float32'),
+ 'double_col': [i * 2.5 for i in range(50)],
+ 'string_col': [f"str{i}" for i in range(50)],
+ 'blob_col': [f"blob{i}".encode('utf-8') for i in range(50)]
+ })
+ writer.write_dataframe(df)
+
+ # Verify by reading back
+ df_read = to_dataframe(tsfile_path, table_name="test_table")
+ # Sort both DataFrames by time for comparison
+ df_read = df_read.sort_values('time').reset_index(drop=True)
+ df_sorted = convert_to_nullable_types(df.sort_values('time').reset_index(drop=True))
+ assert df_read.shape == (50, 8)
+ assert df_read["bool_col"].equals(df_sorted["bool_col"])
+ assert df_read["int32_col"].equals(df_sorted["int32_col"])
+ assert df_read["int64_col"].equals(df_sorted["int64_col"])
+ assert np.allclose(df_read["float_col"], df_sorted["float_col"])
+ assert np.allclose(df_read["double_col"], df_sorted["double_col"])
+ assert df_read["string_col"].equals(df_sorted["string_col"])
+ # BLOB comparison
+ for i in range(50):
+ assert df_read["blob_col"].iloc[i] == df_sorted["blob_col"].iloc[i]
+ finally:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+
+def test_write_dataframe_empty():
+ """Test write_dataframe with empty DataFrame."""
+ table = TableSchema("test_table",
+ [ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD)])
+ tsfile_path = "test_write_dataframe_empty.tsfile"
+ try:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+ with TsFileTableWriter(tsfile_path, table) as writer:
+ # Empty DataFrame
+ df = pd.DataFrame({
+ 'time': [],
+ 'value': []
+ })
+ # Should not raise error (current implementation just passes)
+ with pytest.raises(ValueError) as err:
+ writer.write_dataframe(df)
+
+ finally:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
diff --git a/python/tests/test_to_tsfile.py b/python/tests/test_to_tsfile.py
new file mode 100644
index 00000000..0e4adcd2
--- /dev/null
+++ b/python/tests/test_to_tsfile.py
@@ -0,0 +1,376 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+import os
+
+import numpy as np
+import pandas as pd
+import pytest
+from pandas.core.dtypes.common import is_integer_dtype
+
+from tsfile import to_dataframe, ColumnCategory
+from tsfile.utils import to_tsfile
+
+
+def convert_to_nullable_types(df):
+ """
+ Convert DataFrame columns to nullable types to match returned DataFrame from to_dataframe.
+ This handles the fact that returned DataFrames use nullable types (Int64, Float64, etc.)
+ to support Null values.
+ """
+ df = df.copy()
+ for col in df.columns:
+ dtype = df[col].dtype
+ if dtype == 'int64':
+ df[col] = df[col].astype('Int64')
+ elif dtype == 'int32':
+ df[col] = df[col].astype('Int32')
+ elif dtype == 'float64':
+ df[col] = df[col].astype('Float64')
+ elif dtype == 'float32':
+ df[col] = df[col].astype('Float32')
+ elif dtype == 'bool':
+ df[col] = df[col].astype('boolean')
+ return df
+
+
+def test_to_tsfile_basic():
+ """Test basic to_tsfile functionality with time column."""
+ tsfile_path = "test_to_tsfile_basic.tsfile"
+ try:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+ # Create DataFrame with 'time' column
+ df = pd.DataFrame({
+ 'time': [i for i in range(100)],
+ 'device': [f"device{i}" for i in range(100)],
+ 'value': [i * 1.5 for i in range(100)],
+ 'value2': [i * 10 for i in range(100)]
+ })
+
+ to_tsfile(df, tsfile_path, table_name="test_table")
+
+ # Verify by reading back
+ df_read = to_dataframe(tsfile_path, table_name="test_table")
+ df_read = df_read.sort_values('time').reset_index(drop=True)
+ df_sorted = convert_to_nullable_types(df.sort_values('time').reset_index(drop=True))
+
+ assert df_read.shape == (100, 4)
+ assert df_read["time"].equals(df_sorted["time"])
+ assert df_read["device"].equals(df_sorted["device"])
+ assert df_read["value"].equals(df_sorted["value"])
+ assert df_read["value2"].equals(df_sorted["value2"])
+ finally:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+
+def test_to_tsfile_with_index():
+ """Test to_tsfile using DataFrame index as time when no 'time' column exists."""
+ tsfile_path = "test_to_tsfile_index.tsfile"
+ try:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+ # Create DataFrame without 'time' column, use index
+ df = pd.DataFrame({
+ 'device': [f"device{i}" for i in range(50)],
+ 'value': [i * 2.5 for i in range(50)]
+ })
+ df.index = [i * 10 for i in range(50)] # Set index as timestamps
+
+ to_tsfile(df, tsfile_path, table_name="test_table")
+
+ # Verify by reading back
+ df_read = to_dataframe(tsfile_path, table_name="test_table")
+ df_read = df_read.sort_values('time').reset_index(drop=True)
+ df_sorted = df.sort_index()
+ df_sorted = convert_to_nullable_types(df_sorted.reset_index(drop=True))
+ time_series = pd.Series(df.sort_index().index.values, dtype='Int64')
+
+ assert df_read.shape == (50, 3)
+ assert df_read["time"].equals(time_series)
+ assert df_read["device"].equals(df_sorted["device"])
+ assert df_read["value"].equals(df_sorted["value"])
+ finally:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+
+def test_to_tsfile_custom_time_column():
+ """Test to_tsfile with custom time column name."""
+ tsfile_path = "test_to_tsfile_custom_time.tsfile"
+ try:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+ df = pd.DataFrame({
+ 'timestamp': [i for i in range(30)],
+ 'device': [f"device{i}" for i in range(30)],
+ 'value': [i * 3.0 for i in range(30)]
+ })
+
+ to_tsfile(df, tsfile_path, table_name="test_table", time_column="timestamp")
+
+ # Verify by reading back
+ df_read = to_dataframe(tsfile_path, table_name="test_table")
+ df_read = df_read.sort_values('time').reset_index(drop=True)
+ df_sorted = convert_to_nullable_types(df.sort_values('timestamp').reset_index(drop=True))
+
+ assert df_read.shape == (30, 3)
+ assert df_read["time"].equals(df_sorted["timestamp"])
+ assert df_read["device"].equals(df_sorted["device"])
+ assert df_read["value"].equals(df_sorted["value"])
+ finally:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+
+def test_to_tsfile_with_tag_columns():
+ """Test to_tsfile with tag columns specified."""
+ tsfile_path = "test_to_tsfile_tags.tsfile"
+ try:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+ df = pd.DataFrame({
+ 'time': [i for i in range(20)],
+ 'device': [f"device{i}" for i in range(20)],
+ 'location': [f"loc{i % 5}" for i in range(20)],
+ 'value': [i * 1.5 for i in range(20)]
+ })
+
+ to_tsfile(df, tsfile_path, table_name="test_table", tag_column=["device", "location"])
+
+ # Verify by reading back
+ df_read = to_dataframe(tsfile_path, table_name="test_table")
+ df_read = df_read.sort_values('time').reset_index(drop=True)
+ df_sorted = convert_to_nullable_types(df.sort_values('time').reset_index(drop=True))
+
+ assert df_read.shape == (20, 4)
+ assert df_read["device"].equals(df_sorted["device"])
+ assert df_read["location"].equals(df_sorted["location"])
+ assert df_read["value"].equals(df_sorted["value"])
+ finally:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+
+def test_to_tsfile_all_datatypes():
+ """Test to_tsfile with all supported data types."""
+ tsfile_path = "test_to_tsfile_all_types.tsfile"
+ try:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+ df = pd.DataFrame({
+ 'time': [i for i in range(50)],
+ 'bool_col': [i % 2 == 0 for i in range(50)],
+ 'int32_col': pd.Series([i for i in range(50)], dtype='int32'),
+ 'int64_col': [i * 10 for i in range(50)],
+ 'float_col': pd.Series([i * 1.5 for i in range(50)], dtype='float32'),
+ 'double_col': [i * 2.5 for i in range(50)],
+ 'string_col': [f"str{i}" for i in range(50)],
+ 'blob_col': [f"blob{i}".encode('utf-8') for i in range(50)]
+ })
+
+ to_tsfile(df, tsfile_path, table_name="test_table")
+
+ # Verify by reading back
+ df_read = to_dataframe(tsfile_path, table_name="test_table")
+ df_read = df_read.sort_values('time').reset_index(drop=True)
+ df_sorted = convert_to_nullable_types(df.sort_values('time').reset_index(drop=True))
+
+ assert df_read.shape == (50, 8)
+ assert df_read["bool_col"].equals(df_sorted["bool_col"])
+ assert df_read["int32_col"].equals(df_sorted["int32_col"])
+ assert df_read["int64_col"].equals(df_sorted["int64_col"])
+ assert np.allclose(df_read["float_col"], df_sorted["float_col"])
+ assert np.allclose(df_read["double_col"], df_sorted["double_col"])
+ assert df_read["string_col"].equals(df_sorted["string_col"])
+ # BLOB comparison
+ for i in range(50):
+ assert df_read["blob_col"].iloc[i] == df_sorted["blob_col"].iloc[i]
+ finally:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+
+def test_to_tsfile_default_table_name():
+ """Test to_tsfile with default table name."""
+ tsfile_path = "test_to_tsfile_default_name.tsfile"
+ try:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+ df = pd.DataFrame({
+ 'time': [i for i in range(10)],
+ 'value': [i * 1.0 for i in range(10)]
+ })
+
+ to_tsfile(df, tsfile_path) # No table_name specified
+
+ # Verify by reading back with default table name
+ df_read = to_dataframe(tsfile_path, table_name="table")
+ assert df_read.shape == (10, 2)
+ finally:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+
+def test_to_tsfile_case_insensitive_time():
+ """Test to_tsfile with case-insensitive time column."""
+ tsfile_path = "test_to_tsfile_case_time.tsfile"
+ try:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+ df = pd.DataFrame({
+ 'Time': [i for i in range(20)], # Capital T
+ 'value': [i * 2.0 for i in range(20)]
+ })
+
+ to_tsfile(df, tsfile_path, table_name="test_table")
+
+ # Verify by reading back
+ df_read = to_dataframe(tsfile_path, table_name="test_table")
+ assert df_read.shape == (20, 2)
+ assert df_read["time"].equals(pd.Series([i for i in range(20)], dtype='Int64'))
+ finally:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+
+def test_to_tsfile_empty_dataframe():
+ """Test to_tsfile raises error for empty DataFrame."""
+ tsfile_path = "test_to_tsfile_empty.tsfile"
+ try:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+ df = pd.DataFrame()
+
+ with pytest.raises(ValueError, match="DataFrame cannot be None or empty"):
+ to_tsfile(df, tsfile_path)
+ finally:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+
+def test_to_tsfile_no_data_columns():
+ """Test to_tsfile raises error when only time column exists."""
+ tsfile_path = "test_to_tsfile_no_data.tsfile"
+ try:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+ df = pd.DataFrame({
+ 'time': [i for i in range(10)]
+ })
+
+ with pytest.raises(ValueError, match="DataFrame must have at least one data column"):
+ to_tsfile(df, tsfile_path)
+ finally:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+
+def test_to_tsfile_invalid_time_column():
+ """Test to_tsfile raises error for invalid time column."""
+ tsfile_path = "test_to_tsfile_invalid_time.tsfile"
+ try:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+ df = pd.DataFrame({
+ 'timestamp': [i for i in range(10)],
+ 'value': [i * 1.0 for i in range(10)]
+ })
+
+ # Time column doesn't exist
+ with pytest.raises(ValueError, match="Time column 'time' not found"):
+ to_tsfile(df, tsfile_path, time_column="time")
+ finally:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+
+def test_to_tsfile_non_integer_time_column():
+ """Test to_tsfile raises error for non-integer time column."""
+ tsfile_path = "test_to_tsfile_non_int_time.tsfile"
+ try:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+ df = pd.DataFrame({
+ 'time': [f"time{i}" for i in range(10)], # String, not integer
+ 'value': [i * 1.0 for i in range(10)]
+ })
+
+ with pytest.raises(TypeError, match="must be integer type"):
+ to_tsfile(df, tsfile_path)
+ finally:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+
+def test_to_tsfile_invalid_tag_column():
+ """Test to_tsfile raises error for invalid tag column."""
+ tsfile_path = "test_to_tsfile_invalid_tag.tsfile"
+ try:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+ df = pd.DataFrame({
+ 'time': [i for i in range(10)],
+ 'value': [i * 1.0 for i in range(10)]
+ })
+
+ with pytest.raises(ValueError, match="Tag column 'invalid' not found"):
+ to_tsfile(df, tsfile_path, tag_column=["invalid"])
+ finally:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+
+def test_to_tsfile_string_vs_blob():
+ """Test to_tsfile correctly distinguishes between STRING and BLOB."""
+ tsfile_path = "test_to_tsfile_string_blob.tsfile"
+ try:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
+
+ df = pd.DataFrame({
+ 'time': [i for i in range(20)],
+ 'string_col': [f"str{i}" for i in range(20)], # String
+ 'blob_col': [f"blob{i}".encode('utf-8') for i in range(20)] # Bytes
+ })
+
+ to_tsfile(df, tsfile_path, table_name="test_table")
+
+ # Verify by reading back
+ df_read = to_dataframe(tsfile_path, table_name="test_table")
+ df_read = df_read.sort_values('time').reset_index(drop=True)
+ df_sorted = convert_to_nullable_types(df.sort_values('time').reset_index(drop=True))
+
+ assert df_read["string_col"].equals(df_sorted["string_col"])
+ for i in range(20):
+ assert df_read["blob_col"].iloc[i] == df_sorted["blob_col"].iloc[i]
+ finally:
+ if os.path.exists(tsfile_path):
+ os.remove(tsfile_path)
diff --git a/python/tests/test_write_and_read.py b/python/tests/test_write_and_read.py
index ad8d698a..1ffc22b9 100644
--- a/python/tests/test_write_and_read.py
+++ b/python/tests/test_write_and_read.py
@@ -22,6 +22,7 @@ from datetime import date
import numpy as np
import pandas as pd
import pytest
+from pandas import Float64Dtype
from pandas.core.dtypes.common import is_integer_dtype
from tsfile import ColumnSchema, TableSchema, TSEncoding
@@ -545,7 +546,7 @@ def test_tsfile_to_df():
assert df1.shape == (4097, 4)
assert df1["value2"].sum() == 100 * (1 + 4096) / 2 * 4096
assert is_integer_dtype(df1["time"])
- assert df1["value"].dtype == np.float64
+ assert df1["value"].dtype == Float64Dtype()
assert is_integer_dtype(df1["value2"])
df2 = to_dataframe("table_write_to_df.tsfile", column_names=["device", "value2"])
assert df2.shape == (4097, 3)
@@ -756,457 +757,9 @@ def test_tree_all_datatype_query_to_dataframe_variants():
pass
finally:
- if os.path.exists(tsfile_path):
- os.remove(tsfile_path)
-
-
-def test_table_all_datatype_query_to_dataframe_variants():
- tsfile_path = "test_table.tsfile"
- table = TableSchema(
- "test_table",
- [
- ColumnSchema("Device1", TSDataType.STRING, ColumnCategory.TAG),
- ColumnSchema("Device2", TSDataType.STRING, ColumnCategory.TAG),
- ColumnSchema("Value1", TSDataType.BOOLEAN, ColumnCategory.FIELD),
- ColumnSchema("Value2", TSDataType.INT32, ColumnCategory.FIELD),
- ColumnSchema("Value3", TSDataType.INT64, ColumnCategory.FIELD),
- ColumnSchema("Value4", TSDataType.FLOAT, ColumnCategory.FIELD),
- ColumnSchema("Value5", TSDataType.DOUBLE, ColumnCategory.FIELD),
- ColumnSchema("Value6", TSDataType.TEXT, ColumnCategory.FIELD),
- ColumnSchema("Value7", TSDataType.STRING, ColumnCategory.FIELD),
- ColumnSchema("Value8", TSDataType.BLOB, ColumnCategory.FIELD),
- ColumnSchema("Value9", TSDataType.TIMESTAMP, ColumnCategory.FIELD),
- ColumnSchema("Value10", TSDataType.DATE, ColumnCategory.FIELD),
- ],
- )
- dateSet = set()
- try:
- if os.path.exists(tsfile_path):
- os.remove(tsfile_path)
- max_row_num = 100
- with TsFileTableWriter(tsfile_path, table) as writer:
- tablet = Tablet(
- [
- "Device1",
- "Device2",
- "Value1",
- "Value2",
- "Value3",
- "Value4",
- "Value5",
- "Value6",
- "Value7",
- "Value8",
- "Value9",
- "Value10",
- ],
- [
- TSDataType.STRING,
- TSDataType.STRING,
- TSDataType.BOOLEAN,
- TSDataType.INT32,
- TSDataType.INT64,
- TSDataType.FLOAT,
- TSDataType.DOUBLE,
- TSDataType.TEXT,
- TSDataType.STRING,
- TSDataType.BLOB,
- TSDataType.TIMESTAMP,
- TSDataType.DATE,
- ],
- max_row_num,
- )
- for i in range(max_row_num):
- tablet.add_timestamp(i, i)
- tablet.add_value_by_name("Device1", i, "d1_" + str(i))
- tablet.add_value_by_name("Device2", i, "d2_" + str(i))
- tablet.add_value_by_name("Value1", i, i % 2 == 0)
- tablet.add_value_by_name("Value2", i, i * 3)
- tablet.add_value_by_name("Value3", i, i * 4)
- tablet.add_value_by_name("Value4", i, i * 5.5)
- tablet.add_value_by_name("Value5", i, i * 6.6)
- tablet.add_value_by_name("Value6", i, f"string_value_{i}")
- tablet.add_value_by_name("Value7", i, f"text_value_{i}")
- tablet.add_value_by_name("Value8", i, f"blob_data_{i}".encode('utf-8'))
- tablet.add_value_by_name("Value9", i, i * 9)
- tablet.add_value_by_name("Value10", i, date(2025, 1, i % 20 + 1))
- dateSet.add(date(2025, 1, i % 20 + 1))
- writer.write_table(tablet)
-
- df1_1 = to_dataframe(tsfile_path)
- assert df1_1.shape[0] == max_row_num
- for i in range(max_row_num):
- assert df1_1.iloc[i, 1] == "d1_" + str(df1_1.iloc[i, 0])
- assert df1_1.iloc[i, 2] == "d2_" + str(df1_1.iloc[i, 0])
-
- df2_1 = to_dataframe(tsfile_path, column_names=["Value1"])
- for i in range(max_row_num):
- assert df2_1.iloc[i, 1] == np.bool_(df2_1.iloc[i, 0] % 2 == 0)
- df2_2 = to_dataframe(tsfile_path, column_names=["Value2"])
- for i in range(max_row_num):
- assert df2_2.iloc[i, 1] == np.int32(df2_2.iloc[i, 0] * 3)
- df2_3 = to_dataframe(tsfile_path, column_names=["Value3"])
- for i in range(max_row_num):
- assert df2_3.iloc[i, 1] == np.int64(df2_3.iloc[i, 0] * 4)
- df2_4 = to_dataframe(tsfile_path, column_names=["Value4"])
- for i in range(max_row_num):
- assert df2_4.iloc[i, 1] == np.float32(df2_4.iloc[i, 0] * 5.5)
- df2_5 = to_dataframe(tsfile_path, column_names=["Value5"])
- for i in range(max_row_num):
- assert df2_5.iloc[i, 1] == np.float64(df2_5.iloc[i, 0] * 6.6)
- df2_6 = to_dataframe(tsfile_path, column_names=["Value6"])
- for i in range(max_row_num):
- assert df2_6.iloc[i, 1] == f"string_value_{df2_6.iloc[i, 0]}"
- df2_7 = to_dataframe(tsfile_path, column_names=["Value7"])
- for i in range(max_row_num):
- assert df2_7.iloc[i, 1] == f"text_value_{df2_7.iloc[i, 0]}"
- df2_8 = to_dataframe(tsfile_path, column_names=["Value8"])
- for i in range(max_row_num):
- assert df2_8.iloc[i, 1] == f"blob_data_{df2_8.iloc[i, 0]}".encode('utf-8')
- df2_9 = to_dataframe(tsfile_path, column_names=["Value9"])
- for i in range(max_row_num):
- assert df2_9.iloc[i, 1] == np.int64(df2_9.iloc[i, 0] * 9)
- df2_10 = to_dataframe(tsfile_path, column_names=["Value10"])
- for i in range(max_row_num):
- assert df2_10.iloc[i, 1] in dateSet
- df2_11 = to_dataframe(tsfile_path, column_names=["Device1", "Value1"])
- for i in range(max_row_num):
- assert df2_11.iloc[i, 1] == "d1_" + str(df2_11.iloc[i, 0])
- assert df2_11.iloc[i, 2] == np.bool_(df2_11.iloc[i, 0] % 2 == 0)
- df2_12 = to_dataframe(
- tsfile_path,
- column_names=[
- "Device1",
- "Device2",
- "Value1",
- "Value2",
- "Value3",
- "Value4",
- "Value5",
- "Value6",
- "Value7",
- "Value8",
- "Value9",
- "Value10",
- ],
- )
- for i in range(max_row_num):
- assert df2_12.iloc[i, 1] == "d1_" + str(df2_12.iloc[i, 0])
- assert df2_12.iloc[i, 2] == "d2_" + str(df2_12.iloc[i, 0])
- assert df2_12.iloc[i, 3] == np.bool_(df2_12.iloc[i, 0] % 2 == 0)
- assert df2_12.iloc[i, 4] == np.int32(df2_12.iloc[i, 0] * 3)
- assert df2_12.iloc[i, 5] == np.int64(df2_12.iloc[i, 0] * 4)
- assert df2_12.iloc[i, 6] == np.float32(df2_12.iloc[i, 0] * 5.5)
- assert df2_12.iloc[i, 7] == np.float64(df2_12.iloc[i, 0] * 6.6)
- assert df2_12.iloc[i, 8] == f"string_value_{df2_12.iloc[i, 0]}"
- assert df2_12.iloc[i, 9] == f"text_value_{df2_12.iloc[i, 0]}"
- assert df2_12.iloc[i, 10] == f"blob_data_{df2_12.iloc[i, 0]}".encode(
- "utf-8"
- )
- assert df2_12.iloc[i, 11] == np.int64(df2_12.iloc[i, 0] * 9)
- assert df2_12.iloc[i, 12] == date(2025, 1, df2_12.iloc[i, 0] % 20 + 1)
- df2_13 = to_dataframe(
- tsfile_path, column_names=["Device1", "Device2", "Value1"]
- )
- for i in range(max_row_num):
- assert df2_13.iloc[i, 1] == "d1_" + str(df2_13.iloc[i, 0])
- assert df2_13.iloc[i, 2] == "d2_" + str(df2_13.iloc[i, 0])
- assert df2_13.iloc[i, 3] == np.bool_(df2_13.iloc[i, 0] % 2 == 0)
-
- df3_1 = to_dataframe(tsfile_path, table_name="test_table")
- assert df3_1.shape[0] == max_row_num
- assert df3_1.iloc[0, 0] == 0
- df3_2 = to_dataframe(tsfile_path, table_name="TEST_TABLE")
- assert df3_2.shape[0] == max_row_num
- assert df3_2.iloc[0, 0] == 0
-
- df4_1 = to_dataframe(tsfile_path, start_time=10)
- assert df4_1.shape[0] == 90
- df4_2 = to_dataframe(tsfile_path, start_time=-10)
- assert df4_2.shape[0] == max_row_num
- df4_3 = to_dataframe(tsfile_path, end_time=5)
- assert df4_3.shape[0] == 6
- df4_4 = to_dataframe(tsfile_path, end_time=-5)
- assert df4_4.shape[0] == 0
- df4_5 = to_dataframe(tsfile_path, start_time=5, end_time=5)
- assert df4_5.shape[0] == 1
- df4_6 = to_dataframe(tsfile_path, start_time=-5, end_time=-5)
- assert df4_6.shape[0] == 0
- df4_7 = to_dataframe(tsfile_path, start_time=10, end_time=-10)
- assert df4_7.shape[0] == 0
- df4_8 = to_dataframe(tsfile_path, start_time=-10, end_time=10)
- assert df4_8.shape[0] == 11
- df4_8 = to_dataframe(tsfile_path, start_time=-50, end_time=50)
- assert df4_8.shape[0] == 51
-
- df5_1 = to_dataframe(tsfile_path, max_row_num=1)
- assert df5_1.shape[0] == 1
- df5_2 = to_dataframe(tsfile_path, max_row_num=50)
- assert df5_2.shape[0] == 50
- df5_3 = to_dataframe(tsfile_path, max_row_num=100)
- assert df5_3.shape[0] == 100
- df5_4 = to_dataframe(tsfile_path, max_row_num=1000)
- assert df5_4.shape[0] == 100
- df5_5 = to_dataframe(tsfile_path, max_row_num=0)
- assert df5_5.shape[0] == 0
- df5_6 = to_dataframe(tsfile_path, max_row_num=-10)
- assert df5_6.shape[0] == 0
-
- for df6_1 in to_dataframe(tsfile_path, max_row_num=20, as_iterator=True):
- assert df6_1.shape[0] == 20
- for df6_2 in to_dataframe(tsfile_path, max_row_num=1000, as_iterator=True):
- assert df6_2.shape[0] == 100
-
- for df7_1 in to_dataframe(
- tsfile_path,
- table_name="test_table",
- column_names=["Device1", "Value1"],
- start_time=21,
- end_time=50,
- max_row_num=10,
- as_iterator=True,
- ):
- assert df7_1.shape[0] == 10
- for i in range(30):
- assert df2_11.iloc[i, 1] == "d1_" + str(df2_11.iloc[i, 0])
- assert df2_11.iloc[i, 2] == np.bool_(df2_11.iloc[i, 0] % 2 == 0)
-
- try:
- to_dataframe(tsfile_path, table_name="non_existent_table")
- except TableNotExistError as e:
- assert e.args[0] == "[non_existent_table] Requested table does not exist"
-
- try:
- to_dataframe(tsfile_path, column_names=["non_existent_column"])
- except ColumnNotExistError as e:
- assert e.args[0] == "[non_existent_column] Column does not exist"
-
- finally:
- if os.path.exists(tsfile_path):
- os.remove(tsfile_path)
-
-
-def convert_to_nullable_types(df):
- """
- Convert DataFrame columns to nullable types to match returned DataFrame from to_dataframe.
- This handles the fact that returned DataFrames use nullable types (Int64, Float64, etc.)
- to support Null values.
- """
- df = df.copy()
- for col in df.columns:
- dtype = df[col].dtype
- if dtype == 'int64':
- df[col] = df[col].astype('Int64')
- elif dtype == 'int32':
- df[col] = df[col].astype('Int32')
- elif dtype == 'float64':
- df[col] = df[col].astype('Float64')
- elif dtype == 'float32':
- df[col] = df[col].astype('Float32')
- elif dtype == 'bool':
- df[col] = df[col].astype('boolean')
- return df
-
-
-def test_write_dataframe_basic():
- table = TableSchema("test_table",
- [ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG),
- ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD),
- ColumnSchema("value2", TSDataType.INT64, ColumnCategory.FIELD)])
- tsfile_path = "test_write_dataframe_basic.tsfile"
- try:
- if os.path.exists(tsfile_path):
- os.remove(tsfile_path)
-
- with TsFileTableWriter(tsfile_path, table) as writer:
- # Create DataFrame with 'time' column
- df = pd.DataFrame({
- 'time': [i for i in range(100)],
- 'device': [f"device{i}" for i in range(100)],
- 'value': [i * 1.5 for i in range(100)],
- 'value2': [i * 10 for i in range(100)]
- })
- writer.write_dataframe(df)
-
- df_read = to_dataframe(tsfile_path, table_name="test_table")
- # Sort both DataFrames by time for comparison
- df_read = df_read.sort_values('time').reset_index(drop=True)
- df_sorted = convert_to_nullable_types(df.sort_values('time').reset_index(drop=True))
- assert df_read.shape == (100, 4)
- assert df_read["time"].equals(df_sorted["time"])
- assert df_read["device"].equals(df_sorted["device"])
- assert df_read["value"].equals(df_sorted["value"])
- assert df_read["value2"].equals(df_sorted["value2"])
- finally:
- if os.path.exists(tsfile_path):
- os.remove(tsfile_path)
-
-
-def test_write_dataframe_with_index():
- """Test write_dataframe using DataFrame index as time when no 'time' column exists."""
- table = TableSchema("test_table",
- [ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG),
- ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD)])
- tsfile_path = "test_write_dataframe_index.tsfile"
- try:
- if os.path.exists(tsfile_path):
- os.remove(tsfile_path)
-
- with TsFileTableWriter(tsfile_path, table) as writer:
- # Create DataFrame without 'time' column, use index
- df = pd.DataFrame({
- 'device': [f"device{i}" for i in range(50)],
- 'value': [i * 2.5 for i in range(50)]
- })
- df.index = [i * 10 for i in range(50)] # Set index as timestamps
- writer.write_dataframe(df)
-
- # Verify by reading back
- df_read = to_dataframe(tsfile_path, table_name="test_table")
- # Sort both DataFrames by time for comparison
- df_read = df_read.sort_values('time').reset_index(drop=True)
- df_sorted = df.sort_index()
- df_sorted = convert_to_nullable_types(df_sorted.reset_index(drop=True))
- # Create time series from original index
- time_series = pd.Series(df.sort_index().index.values, dtype='Int64')
- assert df_read.shape == (50, 3)
- # Compare time column with sorted index
- assert df_read["time"].equals(time_series)
- assert df_read["device"].equals(df_sorted["device"])
- assert df_read["value"].equals(df_sorted["value"])
- finally:
- if os.path.exists(tsfile_path):
- os.remove(tsfile_path)
-
-
-def test_write_dataframe_case_insensitive():
- """Test write_dataframe with case-insensitive column matching."""
- table = TableSchema("test_table",
- [ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG),
- ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD)])
- tsfile_path = "test_write_dataframe_case.tsfile"
- try:
- if os.path.exists(tsfile_path):
- os.remove(tsfile_path)
-
- with TsFileTableWriter(tsfile_path, table) as writer:
- # DataFrame with different case column names
- df = pd.DataFrame({
- 'Time': [i for i in range(30)], # Capital T
- 'Device': [f"device{i}" for i in range(30)], # Capital D
- 'VALUE': [i * 3.0 for i in range(30)] # All caps
- })
- writer.write_dataframe(df)
-
- # Verify by reading back
- df_read = to_dataframe(tsfile_path, table_name="test_table")
- # Sort both DataFrames by time for comparison
- df_read = df_read.sort_values('time').reset_index(drop=True)
- df_sorted = convert_to_nullable_types(df.sort_values('Time').reset_index(drop=True))
- assert df_read.shape == (30, 3)
- assert df_read["time"].equals(df_sorted["Time"])
- assert df_read["device"].equals(df_sorted["Device"])
- assert df_read["value"].equals(df_sorted["VALUE"])
- finally:
- if os.path.exists(tsfile_path):
- os.remove(tsfile_path)
-
-
-def test_write_dataframe_column_not_in_schema():
- """Test write_dataframe raises ColumnNotExistError when DataFrame has extra columns."""
- table = TableSchema("test_table",
- [ColumnSchema("device", TSDataType.STRING, ColumnCategory.TAG),
- ColumnSchema("value", TSDataType.DOUBLE, ColumnCategory.FIELD)])
- tsfile_path = "test_write_dataframe_extra_col.tsfile"
- try:
- if os.path.exists(tsfile_path):
- os.remove(tsfile_path)
-
- with TsFileTableWriter(tsfile_path, table) as writer:
- # DataFrame with extra column not in schema
- df = pd.DataFrame({
- 'time': [i for i in range(10)],
- 'device': [f"device{i}" for i in range(10)],
- 'value': [i * 1.0 for i in range(10)],
- 'extra_column': [i for i in range(10)] # Not in schema
- })
- with pytest.raises(ColumnNotExistError) as exc_info:
- writer.write_dataframe(df)
- assert "extra_column" in str(exc_info.value)
- finally:
- if os.path.exists(tsfile_path):
- os.remove(tsfile_path)
-
-
-def test_write_dataframe_type_mismatch():
- """Test write_dataframe raises TypeMismatchError when types are incompatible."""
- table = TableSchema("test_table",
- [ColumnSchema("value", TSDataType.STRING, ColumnCategory.FIELD)])
- tsfile_path = "test_write_dataframe_type_mismatch.tsfile"
- try:
- if os.path.exists(tsfile_path):
- os.remove(tsfile_path)
-
- with TsFileTableWriter(tsfile_path, table) as writer:
- # DataFrame with incompatible type (INT64 vs STRING)
- df = pd.DataFrame({
- 'time': [i for i in range(10)],
- 'value': [i for i in range(10)] # INT64, but schema expects STRING
- })
- with pytest.raises(TypeMismatchError) as exc_info:
- writer.write_dataframe(df)
- assert "Type mismatches" in str(exc_info.value)
- finally:
- if os.path.exists(tsfile_path):
- os.remove(tsfile_path)
-
+ if os.path.exists("tablet_write_and_read.tsfile"):
+ os.remove("tablet_write_and_read.tsfile")
-def test_write_dataframe_all_datatypes():
- """Test write_dataframe with all supported data types."""
- table = TableSchema("test_table",
- [ColumnSchema("bool_col", TSDataType.BOOLEAN, ColumnCategory.FIELD),
- ColumnSchema("int32_col", TSDataType.INT32, ColumnCategory.FIELD),
- ColumnSchema("int64_col", TSDataType.INT64, ColumnCategory.FIELD),
- ColumnSchema("float_col", TSDataType.FLOAT, ColumnCategory.FIELD),
- ColumnSchema("double_col", TSDataType.DOUBLE, ColumnCategory.FIELD),
- ColumnSchema("string_col", TSDataType.STRING, ColumnCategory.FIELD),
- ColumnSchema("blob_col", TSDataType.BLOB, ColumnCategory.FIELD)])
- tsfile_path = "test_write_dataframe_all_types.tsfile"
- try:
- if os.path.exists(tsfile_path):
- os.remove(tsfile_path)
-
- with TsFileTableWriter(tsfile_path, table) as writer:
- df = pd.DataFrame({
- 'time': [i for i in range(50)],
- 'bool_col': [i % 2 == 0 for i in range(50)],
- 'int32_col': pd.Series([i for i in range(50)], dtype='int32'),
- 'int64_col': [i * 10 for i in range(50)],
- 'float_col': pd.Series([i * 1.5 for i in range(50)], dtype='float32'),
- 'double_col': [i * 2.5 for i in range(50)],
- 'string_col': [f"str{i}" for i in range(50)],
- 'blob_col': [f"blob{i}".encode('utf-8') for i in range(50)]
- })
- writer.write_dataframe(df)
-
- # Verify by reading back
- df_read = to_dataframe(tsfile_path, table_name="test_table")
- # Sort both DataFrames by time for comparison
- df_read = df_read.sort_values('time').reset_index(drop=True)
- df_sorted = convert_to_nullable_types(df.sort_values('time').reset_index(drop=True))
- assert df_read.shape == (50, 8)
- assert df_read["bool_col"].equals(df_sorted["bool_col"])
- assert df_read["int32_col"].equals(df_sorted["int32_col"])
- assert df_read["int64_col"].equals(df_sorted["int64_col"])
- assert np.allclose(df_read["float_col"], df_sorted["float_col"])
- assert np.allclose(df_read["double_col"], df_sorted["double_col"])
- assert df_read["string_col"].equals(df_sorted["string_col"])
- # BLOB comparison
- for i in range(50):
- assert df_read["blob_col"].iloc[i] == df_sorted["blob_col"].iloc[i]
- finally:
- if os.path.exists(tsfile_path):
- os.remove(tsfile_path)
if __name__ == "__main__":
os.chdir(os.path.dirname(os.path.abspath(__file__)))
diff --git a/python/tsfile/tsfile_table_writer.py b/python/tsfile/tsfile_table_writer.py
index 54f759c9..9a8498ee 100644
--- a/python/tsfile/tsfile_table_writer.py
+++ b/python/tsfile/tsfile_table_writer.py
@@ -63,11 +63,12 @@ class TsFileTableWriter:
Write a pandas DataFrame into table in tsfile.
:param dataframe: pandas DataFrame with 'time' column and data columns matching schema.
:return: no return value.
+ :raise: ValueError if dataframe is None or is empty.
:raise: ColumnNotExistError if DataFrame columns don't match schema.
:raise: TypeMismatchError if DataFrame column types are incompatible with schema.
"""
if dataframe is None or dataframe.empty:
- pass
+ raise ValueError("DataFrame cannot be None or empty")
# Create mapping from lowercase column name to original column name
df_column_name_map = {col.lower(): col for col in dataframe.columns if col.lower() != 'time'}
diff --git a/python/tsfile/utils.py b/python/tsfile/utils.py
index a8ed5601..a19b7e0e 100644
--- a/python/tsfile/utils.py
+++ b/python/tsfile/utils.py
@@ -221,7 +221,7 @@ def to_tsfile(dataframe: pd.DataFrame,
raise ValueError(f"Time column '{time_column}' not found in DataFrame")
# Check if time column is integer type (int64 or int)
if not is_integer_dtype(dataframe[time_column].dtype):
- raise ValueError(f"Time column '{time_column}' must be integer type (int64 or int), got {dataframe[time_column].dtype}")
+ raise TypeError(f"Time column '{time_column}' must be integer type (int64 or int), got {dataframe[time_column].dtype}")
time_col_name = time_column
else:
# Look for 'time' column (case-insensitive)
@@ -231,6 +231,8 @@ def to_tsfile(dataframe: pd.DataFrame,
if is_integer_dtype(dataframe[col].dtype):
time_col_name = col
break
+ else:
+ raise TypeError(f"Time column '{col}' must be integer type (int64 or int), got {dataframe[col].dtype}")
# Get data columns (excluding time column)
data_columns = [col for col in dataframe.columns if col != time_col_name]
@@ -252,6 +254,15 @@ def to_tsfile(dataframe: pd.DataFrame,
# Infer data type from DataFrame column
col_dtype = dataframe[col_name].dtype
ts_data_type = TSDataType.from_pandas_datatype(col_dtype)
+
+ # If inferred type is STRING but dtype is object, check actual data to distinguish STRING vs BLOB
+ if ts_data_type == TSDataType.STRING and (col_dtype == 'object' or str(col_dtype) == "<class 'numpy.object_'>"):
+ column_series = dataframe[col_name]
+ first_valid_idx = column_series.first_valid_index()
+ if first_valid_idx is not None:
+ first_value = column_series[first_valid_idx]
+ if isinstance(first_value, bytes):
+ ts_data_type = TSDataType.BLOB
# Determine category (TAG or FIELD)
if col_name.lower() in tag_columns_lower:
(tsfile) 03/03: fix.
Posted by co...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
colinlee pushed a commit to branch support_dataframe_to_tsfile
in repository https://gitbox.apache.org/repos/asf/tsfile.git
commit b2de41c475415333a0502714c6813a1e2f9dc6de
Author: ColinLee <sh...@163.com>
AuthorDate: Wed Jan 14 22:20:51 2026 +0800
fix.
---
python/tsfile/tsfile_py_cpp.pyx | 19 +++++++++------
python/tsfile/tsfile_table_writer.py | 10 ++------
python/tsfile/utils.py | 46 +++++++++++++-----------------------
3 files changed, 30 insertions(+), 45 deletions(-)
diff --git a/python/tsfile/tsfile_py_cpp.pyx b/python/tsfile/tsfile_py_cpp.pyx
index 84c550f9..ad0ecdc3 100644
--- a/python/tsfile/tsfile_py_cpp.pyx
+++ b/python/tsfile/tsfile_py_cpp.pyx
@@ -310,6 +310,17 @@ cdef Tablet to_c_tablet(object tablet):
cdef TSDataType pandas_dtype_to_ts_data_type(object dtype):
return to_c_data_type(TSDataTypePy.from_pandas_datatype(dtype))
+cdef TSDataType check_string_or_blob(TSDataType ts_data_type, object dtype, object column_series):
+ if ts_data_type == TS_DATATYPE_STRING:
+ dtype_str = str(dtype)
+ if dtype == 'object' or dtype_str == "<class 'numpy.object_'>":
+ first_valid_idx = column_series.first_valid_index()
+ if first_valid_idx is not None:
+ first_value = column_series[first_valid_idx]
+ if isinstance(first_value, bytes):
+ return TS_DATATYPE_BLOB
+ return ts_data_type
+
cdef Tablet dataframe_to_c_tablet(object target_name, object dataframe):
cdef Tablet ctablet
cdef int max_row_num
@@ -356,13 +367,7 @@ cdef Tablet dataframe_to_c_tablet(object target_name, object dataframe):
for col_name in data_columns:
pandas_dtype = dataframe[col_name].dtype
ds_type = pandas_dtype_to_ts_data_type(pandas_dtype)
- if ds_type == TS_DATATYPE_STRING:
- column_series = dataframe[col_name]
- first_valid_idx = column_series.first_valid_index()
- if first_valid_idx is not None:
- first_value = column_series[first_valid_idx]
- if isinstance(first_value, bytes):
- ds_type = TS_DATATYPE_BLOB
+ ds_type = check_string_or_blob(ds_type, pandas_dtype, dataframe[col_name])
column_types_list.append(ds_type)
columns_names = <char**> malloc(sizeof(char *) * column_num)
diff --git a/python/tsfile/tsfile_table_writer.py b/python/tsfile/tsfile_table_writer.py
index 9a8498ee..e4867c21 100644
--- a/python/tsfile/tsfile_table_writer.py
+++ b/python/tsfile/tsfile_table_writer.py
@@ -21,6 +21,7 @@ from tsfile import TableSchema, Tablet, TableNotExistError
from tsfile import TsFileWriter
from tsfile.constants import TSDataType
from tsfile.exceptions import ColumnNotExistError, TypeMismatchError
+from tsfile.utils import check_string_or_blob
class TsFileTableWriter:
@@ -94,14 +95,7 @@ class TsFileTableWriter:
df_dtype = dataframe[df_col_name_original].dtype
df_ts_type = TSDataType.from_pandas_datatype(df_dtype)
-
- if df_ts_type == TSDataType.STRING and (df_dtype == 'object' or str(df_dtype) == "<class 'numpy.object_'>"):
- column_series = dataframe[df_col_name_original]
- first_valid_idx = column_series.first_valid_index()
- if first_valid_idx is not None:
- first_value = column_series[first_valid_idx]
- if isinstance(first_value, bytes):
- df_ts_type = TSDataType.BLOB
+ df_ts_type = check_string_or_blob(df_ts_type, df_dtype, dataframe[df_col_name_original])
schema_col = schema_column_map[col_name]
expected_ts_type = schema_col.get_data_type()
diff --git a/python/tsfile/utils.py b/python/tsfile/utils.py
index a19b7e0e..d2c19428 100644
--- a/python/tsfile/utils.py
+++ b/python/tsfile/utils.py
@@ -27,6 +27,16 @@ from tsfile.tsfile_reader import TsFileReaderPy
from tsfile import ColumnSchema, TableSchema, ColumnCategory, TSDataType, TsFileTableWriter
+def check_string_or_blob(ts_data_type: TSDataType, dtype, column_series: pd.Series) -> TSDataType:
+ if ts_data_type == TSDataType.STRING and (dtype == 'object' or str(dtype) == "<class 'numpy.object_'>"):
+ first_valid_idx = column_series.first_valid_index()
+ if first_valid_idx is not None:
+ first_value = column_series[first_valid_idx]
+ if isinstance(first_value, bytes):
+ return TSDataType.BLOB
+ return ts_data_type
+
+
def to_dataframe(file_path: str,
table_name: Optional[str] = None,
column_names: Optional[list[str]] = None,
@@ -163,7 +173,7 @@ def to_dataframe(file_path: str,
return pd.DataFrame()
-def to_tsfile(dataframe: pd.DataFrame,
+def dataframe_to_tsfile(dataframe: pd.DataFrame,
file_path: str,
table_name: Optional[str] = None,
time_column: Optional[str] = None,
@@ -209,38 +219,30 @@ def to_tsfile(dataframe: pd.DataFrame,
if dataframe is None or dataframe.empty:
raise ValueError("DataFrame cannot be None or empty")
- # Determine table name
if table_name is None:
table_name = "table"
- # Determine time column
time_col_name = None
if time_column is not None:
- # Check if specified time column exists
if time_column not in dataframe.columns:
raise ValueError(f"Time column '{time_column}' not found in DataFrame")
- # Check if time column is integer type (int64 or int)
if not is_integer_dtype(dataframe[time_column].dtype):
raise TypeError(f"Time column '{time_column}' must be integer type (int64 or int), got {dataframe[time_column].dtype}")
time_col_name = time_column
else:
- # Look for 'time' column (case-insensitive)
for col in dataframe.columns:
if col.lower() == 'time':
- # Check if time column is integer type
if is_integer_dtype(dataframe[col].dtype):
time_col_name = col
break
else:
raise TypeError(f"Time column '{col}' must be integer type (int64 or int), got {dataframe[col].dtype}")
- # Get data columns (excluding time column)
data_columns = [col for col in dataframe.columns if col != time_col_name]
if len(data_columns) == 0:
raise ValueError("DataFrame must have at least one data column besides the time column")
- # Normalize tag_column list (convert to lowercase for comparison)
tag_columns_lower = []
if tag_column is not None:
for tag_col in tag_column:
@@ -248,23 +250,12 @@ def to_tsfile(dataframe: pd.DataFrame,
raise ValueError(f"Tag column '{tag_col}' not found in DataFrame")
tag_columns_lower.append(tag_col.lower())
- # Infer schema from DataFrame columns
column_schemas = []
for col_name in data_columns:
- # Infer data type from DataFrame column
col_dtype = dataframe[col_name].dtype
ts_data_type = TSDataType.from_pandas_datatype(col_dtype)
-
- # If inferred type is STRING but dtype is object, check actual data to distinguish STRING vs BLOB
- if ts_data_type == TSDataType.STRING and (col_dtype == 'object' or str(col_dtype) == "<class 'numpy.object_'>"):
- column_series = dataframe[col_name]
- first_valid_idx = column_series.first_valid_index()
- if first_valid_idx is not None:
- first_value = column_series[first_valid_idx]
- if isinstance(first_value, bytes):
- ts_data_type = TSDataType.BLOB
-
- # Determine category (TAG or FIELD)
+ ts_data_type = check_string_or_blob(ts_data_type, col_dtype, dataframe[col_name])
+
if col_name.lower() in tag_columns_lower:
category = ColumnCategory.TAG
else:
@@ -272,17 +263,12 @@ def to_tsfile(dataframe: pd.DataFrame,
column_schemas.append(ColumnSchema(col_name, ts_data_type, category))
- # Create table schema
table_schema = TableSchema(table_name, column_schemas)
- # Prepare DataFrame for writing
- # If time column is specified or found, rename it to 'time' if needed
- # If no time column is found, write_dataframe will handle it by using index
- df_to_write = dataframe.copy()
if time_col_name is not None and time_col_name != 'time':
- # Rename time column to 'time' for consistency
- df_to_write = df_to_write.rename(columns={time_col_name: 'time'})
+ df_to_write = dataframe.rename(columns={time_col_name: 'time'})
+ else:
+ df_to_write = dataframe
- # Write DataFrame to TsFile
with TsFileTableWriter(file_path, table_schema) as writer:
writer.write_dataframe(df_to_write)
\ No newline at end of file