test: added test cases for more data types

Author: wudidapaopao

programs/local/PandasAnalyzer.cpp

@@ -1,6 +1,7 @@
 #include "PandasAnalyzer.h"
 #include "PythonConversion.h"
 #include "PythonImporter.h"
+#include "base/defines.h"
 
 #include <Common/Exception.h>
 #include <DataTypes/DataTypesNumber.h>
@@ -25,6 +26,12 @@ bool PandasAnalyzer::Analyze(py::object column) {
 	if (sample_size == 0)
 		return false;
 
+	if (sample_size < 0)
+	{
+		analyzed_type = std::make_shared<DataTypeObject>(DataTypeObject::SchemaFormat::JSON);
+		return true;
+	}
+
 	auto & import_cache = PythonImporter::ImportCache();
 	auto pandas = import_cache.pandas();
 	if (!pandas)
@@ -42,7 +49,7 @@ bool PandasAnalyzer::Analyze(py::object column) {
 
 size_t PandasAnalyzer::getSampleIncrement(size_t rows)
 {
-	auto sample = sample_size;
+	auto sample = static_cast<uint64_t>(sample_size);
 	if (sample > rows)
 		sample = rows;
 

programs/local/PandasAnalyzer.h

@@ -30,7 +30,7 @@ class PandasAnalyzer {
 	size_t getSampleIncrement(size_t rows);
 
 private:
-	uint64_t sample_size;
+	int64_t sample_size;
 	PythonGILWrapper gil;
 	DB::DataTypePtr analyzed_type;
 };

programs/local/PythonConversion.cpp

@@ -90,7 +90,7 @@ bool isInteger(const py::handle & obj)
 	return GetPythonObjectType(obj) == PythonObjectType::Integer;
 }
 
-void writeInteger(const py::handle & obj, rapidjson::Value & json_value)
+static void writeInteger(const py::handle & obj, rapidjson::Value & json_value)
 {
 	auto ptr = obj.ptr();
 	int overflow = 0;
@@ -137,7 +137,7 @@ bool isNone(const py::handle & obj)
 	return GetPythonObjectType(obj) == PythonObjectType::None;
 }
 
-void writeNone(const py::handle & obj, rapidjson::Value & json_value)
+static void writeNone(const py::handle & obj, rapidjson::Value & json_value)
 {
 	json_value.SetNull();
 }
@@ -147,25 +147,25 @@ bool isFloat(const py::handle & obj)
 	return GetPythonObjectType(obj) == PythonObjectType::Float;
 }
 
-void writeFloat(const py::handle & obj, rapidjson::Value & json_value)
+static void writeFloat(const py::handle & obj, rapidjson::Value & json_value)
 {
 	auto ptr = obj.ptr();
-	if (std::isnan(PyFloat_AsDouble(ptr)))
-	{
-		json_value.SetNull();
-		return;
-	}
+    double value = PyFloat_AsDouble(ptr);
 
-	double value = obj.cast<double>();
-	json_value.SetDouble(value);
+    if (std::isnan(value) || std::isinf(value)) {
+        json_value.SetNull();
+        return;
+    }
+
+    json_value.SetDouble(value);
 }
 
 bool isBoolean(const py::handle & obj)
 {
 	return GetPythonObjectType(obj) == PythonObjectType::Bool;
 }
 
-void writeBoolean(const py::handle & obj, rapidjson::Value & json_value)
+static void writeBoolean(const py::handle & obj, rapidjson::Value & json_value)
 {
 	json_value.SetBool(py::cast<bool>(obj));
 }
@@ -175,9 +175,9 @@ bool isDecimal(const py::handle & obj)
 	return GetPythonObjectType(obj) == PythonObjectType::Decimal;
 }
 
-void writeDecimal(const py::handle & obj, rapidjson::Value & json_value, rapidjson::Document::AllocatorType & allocator)
+static void writeDecimal(const py::handle & obj, rapidjson::Value & json_value, rapidjson::Document::AllocatorType & allocator)
 {
-	const auto & str = obj.cast<std::string>();
+    String str = py::str(obj);
 	json_value.SetString(str.data(), str.size(), allocator);
 }
 
@@ -187,18 +187,84 @@ bool isString(const py::handle & obj)
 		|| GetPythonObjectType(obj) == PythonObjectType::Bytes;
 }
 
-void writeString(const py::handle & obj, rapidjson::Value & json_value, rapidjson::Document::AllocatorType & allocator)
+bool isByteArray(const py::handle & obj)
+{
+	return GetPythonObjectType(obj) == PythonObjectType::ByteArray;
+}
+
+static void writeByteArray(const py::handle & obj, rapidjson::Value & json_value, rapidjson::Document::AllocatorType & allocator)
+{
+    auto * ptr = obj.ptr();
+	auto * data = PyByteArray_AsString(ptr);
+    auto size = PyByteArray_GET_SIZE(ptr);
+    json_value.SetString(data, size, allocator);
+}
+
+bool isMemoryView(const py::handle & obj)
+{
+    return GetPythonObjectType(obj) == PythonObjectType::MemoryView;
+}
+
+static void writeMemoryView(const py::handle & obj, rapidjson::Value & json_value, rapidjson::Document::AllocatorType & allocator)
+{
+    auto * ptr = obj.ptr();
+    py::memoryview py_view = obj.cast<py::memoryview>();
+
+	Py_buffer * py_buf = PyMemoryView_GET_BUFFER(ptr);
+    auto * data = static_cast<char *>(py_buf->buf);
+    auto size = py_buf->len;
+
+    json_value.SetString(data, size, allocator);
+}
+
+static void writeString(const py::handle & obj, rapidjson::Value & json_value, rapidjson::Document::AllocatorType & allocator)
 {
 	const auto & str = obj.cast<std::string>();
 	json_value.SetString(str.data(), str.size(), allocator);
 }
 
-void writeOthers(const py::handle & obj, rapidjson::Value & json_value, rapidjson::Document::AllocatorType & allocator)
+static void writeOthers(const py::handle & obj, rapidjson::Value & json_value, rapidjson::Document::AllocatorType & allocator)
 {
 	String str = py::str(obj);
 	json_value.SetString(str.data(), str.size(), allocator);
 }
 
+void handlePrimitiveTypes(
+    const py::handle & obj, rapidjson::Value & json_value, rapidjson::Document::AllocatorType & allocator)
+{
+    PythonObjectType type = GetPythonObjectType(obj);
+    switch(type) {
+    case PythonObjectType::Integer:
+        writeInteger(obj, json_value);
+        break;
+    case PythonObjectType::Float:
+        writeFloat(obj, json_value);
+        break;
+    case PythonObjectType::Bool:
+        writeBoolean(obj, json_value);
+        break;
+    case PythonObjectType::Decimal:
+        writeDecimal(obj, json_value, allocator);
+        break;
+    case PythonObjectType::String:
+    case PythonObjectType::Bytes:
+        writeString(obj, json_value, allocator);
+        break;
+    case PythonObjectType::ByteArray:
+        writeByteArray(obj, json_value, allocator);
+        break;
+    case PythonObjectType::MemoryView:
+        writeMemoryView(obj, json_value, allocator);
+        break;
+    case PythonObjectType::None:
+        writeNone(obj, json_value);
+        break;
+    default:
+        writeOthers(obj, json_value, allocator);
+        break;
+    }
+}
+
 void convert_to_json_str(const py::handle & obj, String & ret)
 {
     rapidjson::Document d;
@@ -213,8 +279,8 @@ void convert_to_json_str(const py::handle & obj, String & ret)
             for (auto& item : py::cast<py::dict>(obj))
             {
                 rapidjson::Value key;
-                auto ket_str = py::str(item.first).cast<std::string>();
-                key.SetString(ket_str.data(), ket_str.size(), allocator);
+                auto key_str = py::str(item.first).cast<std::string>();
+                key.SetString(key_str.data(), key_str.size(), allocator);
 
                 rapidjson::Value val;
                 convert(item.second, val);
@@ -259,33 +325,9 @@ void convert_to_json_str(const py::handle & obj, String & ret)
                 json_value.PushBack(element, allocator);
 		    }
         }
-        else if (isInteger(obj))
-        {
-            writeInteger(obj, json_value);
-        }
-        else if (isNone(obj))
-        {
-            writeNone(obj, json_value);
-        }
-        else if (isFloat(obj))
-        {
-            writeFloat(obj, json_value);
-        }
-        else if (isBoolean(obj))
-        {
-            writeBoolean(obj, json_value);
-        }
-        else if (isDecimal(obj))
-        {
-            writeDecimal(obj, json_value, allocator);
-        }
-        else if (isString(obj))
-        {
-            writeString(obj, json_value, allocator);
-        }
         else
         {
-            writeOthers(obj, json_value, allocator);
+            handlePrimitiveTypes(obj, json_value, allocator);
         }
     };
 

programs/local/PythonConversion.h

@@ -28,37 +28,26 @@ enum class PythonObjectType {
 	Tuple,
 	Dict,
 	NdArray,
-	NdDatetime,
-	Value
+	NdDatetime
 };
 
 PythonObjectType GetPythonObjectType(const py::handle & obj);
 
 bool isInteger(const py::handle & obj);
 
-void writeInteger(const py::handle & obj, rapidjson::Value & json_value);
-
 bool isNone(const py::handle & obj);
 
-void writeNone(const py::handle & obj, rapidjson::Value & json_value);
-
 bool isFloat(const py::handle & obj);
 
-void writeFloat(const py::handle & obj, rapidjson::Value & json_value);
-
 bool isBoolean(const py::handle & obj);
 
-void writeBoolean(const py::handle & obj, rapidjson::Value & json_value);
-
 bool isDecimal(const py::handle & obj);
 
-void writeDecimal(const py::handle & obj, rapidjson::Value & json_value, rapidjson::Document::AllocatorType & allocator);
-
 bool isString(const py::handle & obj);
 
-void writeString(const py::handle & obj, rapidjson::Value & json_value, rapidjson::Document::AllocatorType & allocator);
+bool isByteArray(const py::handle & obj);
 
-void writeOthers(const py::handle & obj, rapidjson::Value & json_value, rapidjson::Document::AllocatorType & allocator);
+bool isMemoryView(const py::handle & obj);
 
 void convert_to_json_str(const py::handle & obj, String & ret);
 

src/Core/Settings.h

@@ -968,7 +968,7 @@ class IColumn;
     \
     /** Experimental feature for moving data between shards. */ \
     M(Bool, allow_experimental_query_deduplication, false, "Experimental data deduplication for SELECT queries based on part UUIDs", 0) \
-    M(UInt64, pandas_analyze_sample, 10000, "Sample rows in pandas to automatically determine the data types. When set to 0, sampling is disabled", 0) \
+    M(Int64, pandas_analyze_sample, 10000, "Sample rows in pandas to automatically determine the data types. When set to 0, sampling is disabled", 0) \
 
     /** End of experimental features */