Merge pull request #330 from wudidapaopao/yuxiaozhe-dev

Implement JSON type support

Author: auxten

programs/local/CMakeLists.txt

@@ -1,22 +1,50 @@
 set (CLICKHOUSE_LOCAL_SOURCES LocalServer.cpp)
 
 if (USE_PYTHON)
-    set(CLICKHOUSE_LOCAL_SOURCES ${CLICKHOUSE_LOCAL_SOURCES} LocalChdb.cpp)
+    set (CHDB_SOURCES
+        FormatHelper.cpp
+        ListScan.cpp
+        LocalChdb.cpp
+        LocalServer.cpp
+        NumpyType.cpp
+        PandasAnalyzer.cpp
+        PandasDataFrame.cpp
+        PandasScan.cpp
+        PybindWrapper.cpp
+        PythonConversion.cpp
+        PythonDict.cpp
+        PythonReader.cpp
+        PythonTableCache.cpp
+        PythonImportCache.cpp
+        PythonImporter.cpp
+        PythonSource.cpp
+        PythonUtils.cpp
+        StoragePython.cpp
+        TableFunctionPython.cpp
+    )
+    set (CLICKHOUSE_LOCAL_SOURCES ${CLICKHOUSE_LOCAL_SOURCES} ${CHDB_SOURCES})
+
     # include path from shell cmd "python3 -m pybind11 --includes"
     execute_process(COMMAND python3 -m pybind11 --includes
         OUTPUT_VARIABLE PYBIND11_INCLUDES
         OUTPUT_STRIP_TRAILING_WHITESPACE
     )
     string(REGEX REPLACE ".*-I([^ ]+).*" "\\1" PYBIND11_INCLUDE_DIR ${PYBIND11_INCLUDES})
-    include_directories(${PYBIND11_INCLUDE_DIR})
+    include_directories(${PYBIND11_INCLUDE_DIR} ${CMAKE_CURRENT_SOURCE_DIR})
 
     # include Python.h
     execute_process(COMMAND python3-config --includes
         OUTPUT_VARIABLE PYTHON_INCLUDES
         OUTPUT_STRIP_TRAILING_WHITESPACE
     )
     string(REGEX REPLACE ".*-I([^ ]+).*" "\\1" PYTHON_INCLUDE_DIR ${PYTHON_INCLUDES})
-    set_source_files_properties(LocalChdb.cpp PROPERTIES INCLUDE_DIRECTORIES ${PYTHON_INCLUDE_DIR})
+
+    foreach(_file ${CHDB_SOURCES})
+        set_source_files_properties(${_file}
+            PROPERTIES INCLUDE_DIRECTORIES
+            ${PYTHON_INCLUDE_DIR}
+        )
+    endforeach(_file)
 
     # get python version, something like python3.x
     execute_process(COMMAND python3 -c "import sys; print('python3.'+str(sys.version_info[1]))"
@@ -32,18 +60,27 @@ if (USE_PYTHON)
         #     "-w -idirafter /usr/include -include x86_64-linux-gnu/${PYTHON_VERSION}/pyconfig.h"
         # )
         if (PYTHON_VERSION STREQUAL "python3.6" OR PYTHON_VERSION STREQUAL "python3.7" OR PYTHON_VERSION STREQUAL "python3.8")
-            set_source_files_properties(LocalChdb.cpp PROPERTIES COMPILE_FLAGS
-                "-w -idirafter /usr/include -include crypt.h"
-            )
+            foreach(_file ${CHDB_SOURCES})
+                set_source_files_properties(${_file}
+                    PROPERTIES COMPILE_FLAGS
+                    "-w -idirafter /usr/include -include crypt.h"
+                )
+            endforeach(_file)
         else()
-            set_source_files_properties(LocalChdb.cpp PROPERTIES COMPILE_FLAGS
-                "-w"
-            )
+            foreach(_file ${CHDB_SOURCES})
+                set_source_files_properties(${_file}
+                    PROPERTIES COMPILE_FLAGS
+                    "-w"
+                )
+            endforeach(_file)
         endif()
     elseif (OS_DARWIN)
-        set_source_files_properties(LocalChdb.cpp PROPERTIES COMPILE_FLAGS 
-            "-w"
-        )
+        foreach(_file ${CHDB_SOURCES})
+            set_source_files_properties(${_file}
+                PROPERTIES COMPILE_FLAGS
+                "-w"
+            )
+        endforeach(_file)
     endif()
 endif()
 

programs/local/DatetimeCacheItem.h

@@ -0,0 +1,60 @@
+#pragma once
+
+#include "PythonImportCacheItem.h"
+
+namespace CHDB {
+
+struct DatetimeDatetimeCacheItem : public PythonImportCacheItem
+{
+
+public:
+	DatetimeDatetimeCacheItem(PythonImportCacheItem * parent)
+	    : PythonImportCacheItem("datetime", parent), min("min", this), max("max", this), combine("combine", this)
+	{
+	}
+
+	~DatetimeDatetimeCacheItem() override = default;
+
+	PythonImportCacheItem min;
+	PythonImportCacheItem max;
+	PythonImportCacheItem combine;
+};
+
+struct DatetimeDateCacheItem : public PythonImportCacheItem
+{
+
+public:
+	DatetimeDateCacheItem(PythonImportCacheItem * parent)
+	    : PythonImportCacheItem("date", parent), max("max", this), min("min", this)
+	{
+	}
+
+	~DatetimeDateCacheItem() override = default;
+
+	PythonImportCacheItem max;
+	PythonImportCacheItem min;
+};
+
+struct DatetimeCacheItem : public PythonImportCacheItem
+{
+
+public:
+	static constexpr const char *Name = "datetime";
+
+public:
+	DatetimeCacheItem()
+	    : PythonImportCacheItem("datetime"), date(this), time("time", this), timedelta("timedelta", this),
+	      datetime(this), timezone("timezone", this)
+	{
+	}
+
+	~DatetimeCacheItem() override = default;
+
+	DatetimeDateCacheItem date;
+	PythonImportCacheItem time;
+	PythonImportCacheItem timedelta;
+	DatetimeDatetimeCacheItem datetime;
+	PythonImportCacheItem timezone;
+};
+
+} // namespace CHDB

programs/local/DecimalCacheItem.h

@@ -0,0 +1,21 @@
+#pragma once
+
+#include "PythonImportCacheItem.h"
+
+namespace CHDB {
+
+struct DecimalCacheItem : public PythonImportCacheItem
+{
+public:
+	static constexpr const char * Name = "decimal";
+
+	DecimalCacheItem() : PythonImportCacheItem("decimal"), Decimal("Decimal", this)
+	{
+	}
+
+	~DecimalCacheItem() override = default;
+
+	PythonImportCacheItem Decimal;
+};
+
+} // namespace CHDB

programs/local/FormatHelper.cpp

@@ -0,0 +1,28 @@
+#include "FormatHelper.h"
+
+namespace CHDB {
+
+static bool is_json_supported = true;
+
+void SetCurrentFormat(const char * format)
+{
+    if (format)
+    {
+        String lowerFormat = format;
+        std::transform(lowerFormat.begin(), lowerFormat.end(), lowerFormat.begin(), ::tolower);
+
+        is_json_supported =  !(lowerFormat == "arrow" || lowerFormat == "parquet" || lowerFormat == "arrowstream"
+            || lowerFormat == "protobuf" || lowerFormat == "protobuflist" || lowerFormat == "protobufsingle");
+
+        return;
+    }
+
+    is_json_supported = true;
+}
+
+bool isJSONSupported()
+{
+    return is_json_supported;
+}
+
+} // namespace CHDB

programs/local/FormatHelper.h

@@ -0,0 +1,11 @@
+#pragma once
+
+#include <base/types.h>
+
+namespace CHDB {
+
+void SetCurrentFormat(const char * format);
+
+bool isJSONSupported();
+
+} // namespace CHDB

programs/local/ListScan.cpp

@@ -0,0 +1,83 @@
+#include "ListScan.h"
+#include "PythonConversion.h"
+
+#include <Columns/ColumnObject.h>
+#include <DataTypes/IDataType.h>
+#include <DataTypes/Serializations/SerializationJSON.h>
+#include <IO/WriteHelpers.h>
+
+namespace DB
+{
+
+namespace ErrorCodes
+{
+    extern const int BAD_ARGUMENTS;
+    extern const int PY_EXCEPTION_OCCURED;
+}
+
+}
+
+using namespace DB;
+
+namespace CHDB {
+
+ColumnPtr ListScan::scanObject(
+    const ColumnWrapper & col_wrap,
+    const size_t cursor,
+    const size_t count,
+    const FormatSettings & format_settings)
+{
+    innerCheck(col_wrap);
+
+    auto & data_type = col_wrap.dest_type;
+    auto column = data_type->createColumn();
+    auto serialization = data_type->getDefaultSerialization();
+
+    innerScanObject(cursor, count, format_settings, serialization, col_wrap.data, column);
+
+    return column;
+}
+
+void ListScan::scanObject(
+    const size_t cursor,
+    const size_t count,
+    const FormatSettings & format_settings,
+    const py::handle & obj,
+    MutableColumnPtr & column)
+{
+    auto data_type = std::make_shared<DataTypeObject>(DataTypeObject::SchemaFormat::JSON);
+    SerializationPtr serialization = data_type->getDefaultSerialization();
+
+    innerScanObject(cursor, count, format_settings, serialization, obj, column);
+}
+
+void ListScan::innerScanObject(
+    const size_t cursor,
+    const size_t count,
+    const FormatSettings & format_settings,
+    SerializationPtr & serialization,
+    const py::handle & obj,
+    MutableColumnPtr & column)
+{
+    py::gil_scoped_acquire acquire;
+
+    auto list = obj.cast<py::list>();
+
+    for (size_t i = cursor; i < cursor + count; ++i)
+    {
+        auto item = list.attr("__getitem__")(i);
+        if (!tryInsertJsonResult(item, format_settings, column, serialization))
+            column->insertDefault();
+	}
+}
+
+void ListScan::innerCheck(const ColumnWrapper & col_wrap)
+{
+    if (col_wrap.data.is_none())
+        throw Exception(ErrorCodes::PY_EXCEPTION_OCCURED, "Column data is None");
+
+    if (!col_wrap.buf)
+        throw Exception(ErrorCodes::PY_EXCEPTION_OCCURED, "Column buffer is null");
+}
+
+} // namespace CHDB

programs/local/ListScan.h

@@ -0,0 +1,35 @@
+#pragma once
+
+#include "PybindWrapper.h"
+#include "PythonUtils.h"
+
+namespace CHDB {
+
+class ListScan {
+public:
+    static DB::ColumnPtr scanObject(
+        const DB::ColumnWrapper & col_wrap,
+        const size_t cursor,
+        const size_t count,
+        const DB::FormatSettings & format_settings);
+
+    static void scanObject(
+        const size_t cursor,
+        const size_t count,
+        const DB::FormatSettings & format_settings,
+        const py::handle & obj,
+        DB::MutableColumnPtr & column);
+
+private:
+    static void innerCheck(const DB::ColumnWrapper & col_wrap);
+
+    static void innerScanObject(
+        const size_t cursor,
+        const size_t count,
+        const DB::FormatSettings & format_settings,
+        DB::SerializationPtr & serialization,
+        const py::handle & obj,
+        DB::MutableColumnPtr & column);
+};
+
+} // namespace CHDB