feat: implement and test querying JSON columns in pyarrow Table, DataFrame, dict, and PyReader

- Implemented support for querying JSON columns across various data sources including pyarrow Table, DataFrame, dictionary, and PyReader.
- Added corresponding test cases to validate the querying functionality for each data type.
- Enhanced the display format of run_all.py.

Author: wudidapaopao

programs/local/CMakeLists.txt

@@ -2,6 +2,8 @@ set (CLICKHOUSE_LOCAL_SOURCES LocalServer.cpp)
 
 if (USE_PYTHON)
     set (CHDB_SOURCES
+        FormatHelper.cpp
+        ListScan.cpp
         LocalChdb.cpp
         LocalServer.cpp
         NumpyType.cpp
@@ -10,6 +12,9 @@ if (USE_PYTHON)
         PandasScan.cpp
         PybindWrapper.cpp
         PythonConversion.cpp
+        PythonDict.cpp
+        PythonReader.cpp
+        PythonTableCache.cpp
         PythonImportCache.cpp
         PythonImporter.cpp
         PythonObjects.cpp

programs/local/DatetimeCacheItem.h

@@ -4,7 +4,8 @@
 
 namespace CHDB {
 
-struct DatetimeDatetimeCacheItem : public PythonImportCacheItem {
+struct DatetimeDatetimeCacheItem : public PythonImportCacheItem
+{
 
 public:
 	DatetimeDatetimeCacheItem(PythonImportCacheItem * parent)
@@ -19,7 +20,8 @@ struct DatetimeDatetimeCacheItem : public PythonImportCacheItem {
 	PythonImportCacheItem combine;
 };
 
-struct DatetimeDateCacheItem : public PythonImportCacheItem {
+struct DatetimeDateCacheItem : public PythonImportCacheItem
+{
 
 public:
 	DatetimeDateCacheItem(PythonImportCacheItem * parent)
@@ -33,7 +35,8 @@ struct DatetimeDateCacheItem : public PythonImportCacheItem {
 	PythonImportCacheItem min;
 };
 
-struct DatetimeCacheItem : public PythonImportCacheItem {
+struct DatetimeCacheItem : public PythonImportCacheItem
+{
 
 public:
 	static constexpr const char *Name = "datetime";

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 "DataTypes/IDataType.h"
+#include "IO/WriteHelpers.h"
+#include "PythonConversion.h"
+
+#include <Columns/ColumnObject.h>
+#include <DataTypes/Serializations/SerializationJSON.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

programs/local/LocalChdb.cpp

@@ -2,68 +2,16 @@
 #include "LocalServer.h"
 #include "chdb.h"
 #include "PythonImporter.h"
+#include "PythonTableCache.h"
 #include "TableFunctionPython.h"
 
 #include <mutex>
-#include "Common/logger_useful.h"
-#include <Common/re2.h>
-#include "pybind11/gil.h"
-#include "pybind11/pytypes.h"
+#include <Common/logger_useful.h>
 
 namespace py = pybind11;
 
 extern bool inside_main = true;
 
-// Global storage for Python Table Engine queriable object
-extern py::handle global_query_obj;
-
-// Find the queriable object in the Python environment
-// return nullptr if no Python obj is referenced in query string
-// return py::none if the obj referenced not found
-// return the Python object if found
-// The object name is extracted from the query string, must referenced by
-// Python(var_name) or Python('var_name') or python("var_name") or python('var_name')
-// such as:
-//  - `SELECT * FROM Python('PyReader')`
-//  - `SELECT * FROM Python(PyReader_instance)`
-//  - `SELECT * FROM Python(some_var_with_type_pandas_DataFrame_or_pyarrow_Table)`
-// The object can be any thing that Python Table supported, like PyReader, pandas DataFrame, or PyArrow Table
-// The object should be in the global or local scope
-py::handle findQueryableObjFromQuery(const std::string & query_str)
-{
-    // Extract the object name from the query string
-    std::string var_name;
-
-    // RE2 pattern to match Python()/python() patterns with single/double quotes or no quotes
-    static const RE2 pattern(R"([Pp]ython\s*\(\s*(?:['"]([^'"]+)['"]|([a-zA-Z_][a-zA-Z0-9_]*))\s*\))");
-
-    re2::StringPiece input(query_str);
-    std::string quoted_match, unquoted_match;
-
-    // Try to match and extract the groups
-    if (RE2::PartialMatch(query_str, pattern, &quoted_match, &unquoted_match))
-    {
-        // If quoted string was matched
-        if (!quoted_match.empty())
-        {
-            var_name = quoted_match;
-        }
-        // If unquoted identifier was matched
-        else if (!unquoted_match.empty())
-        {
-            var_name = unquoted_match;
-        }
-    }
-
-    if (var_name.empty())
-    {
-        return nullptr;
-    }
-
-    // Find the object in the Python environment
-    return DB::findQueryableObj(var_name);
-}
-
 local_result_v2 * queryToBuffer(
     const std::string & queryStr,
     const std::string & output_format = "CSV",
@@ -303,7 +251,7 @@ void connection_wrapper::commit()
 
 query_result * connection_wrapper::query(const std::string & query_str, const std::string & format)
 {
-    global_query_obj = findQueryableObjFromQuery(query_str);
+    CHDB::PythonTableCache::findQueryableObjFromQuery(query_str);
 
     py::gil_scoped_release release;
     auto * result = query_conn(*conn, query_str.c_str(), format.c_str());
@@ -322,7 +270,7 @@ query_result * connection_wrapper::query(const std::string & query_str, const st
 
 streaming_query_result * connection_wrapper::send_query(const std::string & query_str, const std::string & format)
 {
-    global_query_obj = findQueryableObjFromQuery(query_str);
+    CHDB::PythonTableCache::findQueryableObjFromQuery(query_str);
 
     py::gil_scoped_release release;
     auto * result = query_conn_streaming(*conn, query_str.c_str(), format.c_str());
@@ -372,7 +320,7 @@ void connection_wrapper::streaming_cancel_query(streaming_query_result * streami
 void cursor_wrapper::execute(const std::string & query_str)
 {
     release_result();
-    global_query_obj = findQueryableObjFromQuery(query_str);
+    CHDB::PythonTableCache::findQueryableObjFromQuery(query_str);
 
     // Use JSONCompactEachRowWithNamesAndTypes format for better type support
     py::gil_scoped_release release;
@@ -545,6 +493,7 @@ PYBIND11_MODULE(_chdb, m)
 	auto destroy_import_cache = []()
     {
         DB::LocalServer::cleanupConnection();
+        CHDB::PythonTableCache::clear();
 		CHDB::PythonImporter::destroy();
 	};
 	m.add_object("_destroy_import_cache", py::capsule(destroy_import_cache));

programs/local/LocalServer.cpp

@@ -3,7 +3,9 @@
 #include "chdb-internal.h"
 
 #if USE_PYTHON
+#include "FormatHelper.h"
 #include "TableFunctionPython.h"
+#include "PythonTableCache.h"
 #include <TableFunctions/TableFunctionFactory.h>
 #include <Storages/StorageFactory.h>
 #endif
@@ -1425,6 +1427,8 @@ static CHDB::ResultData createQueryResult(DB::LocalServer * server, const CHDB::
             /// Acquiring GIL during process termination leads to immediate thread termination.
             local_connection->resetQueryContext();
         }
+
+        CHDB::PythonTableCache::clear();
 #endif
     }
 
@@ -1704,13 +1708,19 @@ static CHDB::ResultData executeQueryRequest(
                 streaming_req->query = query;
                 streaming_req->format = format;
                 queue->current_query = std::move(streaming_req);
+#if USE_PYTHON
+                CHDB::SetCurrentFormat(format);
+#endif
             }
             else if (query_type == CHDB::QueryType::TYPE_MATERIALIZED)
             {
                 auto materialized_req = std::make_unique<CHDB::MaterializedQueryRequest>();
                 materialized_req->query = query;
                 materialized_req->format = format;
                 queue->current_query = std::move(materialized_req);
+#if USE_PYTHON
+                CHDB::SetCurrentFormat(format);
+#endif
             }
             else
             {

programs/local/PandasAnalyzer.cpp

@@ -60,13 +60,9 @@ DataTypePtr PandasAnalyzer::getItemType(py::object obj, bool & can_convert)
 
 	switch (object_type) {
 	case PythonObjectType::Dict:
+		return std::make_shared<DataTypeObject>(DataTypeObject::SchemaFormat::JSON);
 	case PythonObjectType::Tuple:
 	case PythonObjectType::List:
-	{
-		// PyDictionary dict = PyDictionary(py::reinterpret_borrow<py::object>(obj));
-
-		return std::make_shared<DataTypeObject>(DataTypeObject::SchemaFormat::JSON);
-	}
 	case PythonObjectType::None:
 	case PythonObjectType::Bool:
 	case PythonObjectType::Integer: