import chdb
from typing import List, Tuple, Any
import json
class DatabaseReader(chdb.PyReader):
"""Custom reader for database-like data sources"""
def __init__(self, connection_string: str):
# Simulate database connection
self.data = self._load_data(connection_string)
self.cursor = 0
self.batch_size = 1000
super().__init__(self.data)
def _load_data(self, conn_str):
# Simulate loading from database
return {
"id": list(range(1, 10001)),
"name": [f"user_{i}" for i in range(1, 10001)],
"score": [i * 10 + (i % 7) for i in range(1, 10001)],
"metadata": [
json.dumps({"level": i % 5, "active": i % 3 == 0})
for i in range(1, 10001)
]
}
def get_schema(self) -> List[Tuple[str, str]]:
"""Define table schema with explicit types"""
return [
("id", "UInt64"),
("name", "String"),
("score", "Int64"),
("metadata", "String") # JSON stored as string
]
def read(self, col_names: List[str], count: int) -> List[List[Any]]:
"""Read data in batches"""
if self.cursor >= len(self.data["id"]):
return [] # No more data
end_pos = min(self.cursor + min(count, self.batch_size), len(self.data["id"]))
# Return data for requested columns
result = []
for col in col_names:
if col in self.data:
result.append(self.data[col][self.cursor:end_pos])
else:
# Handle missing columns
result.append([None] * (end_pos - self.cursor))
self.cursor = end_pos
return result
### JSON Type Inference and Handling
chDB automatically handles complex nested data structures:
```python
import pandas as pd
import chdb
# DataFrame with mixed JSON objects
df_with_json = pd.DataFrame({
"user_id": [1, 2, 3, 4],
"profile": [
{"name": "Alice", "age": 25, "preferences": ["music", "travel"]},
{"name": "Bob", "age": 30, "location": {"city": "NYC", "country": "US"}},
{"name": "Charlie", "skills": ["python", "sql", "ml"], "experience": 5},
{"score": 95, "rank": "gold", "achievements": [{"title": "Expert", "date": "2024-01-01"}]}
]
})
# Control JSON inference with settings
result = chdb.query("""
SELECT
user_id,
profile.name as name,
profile.age as age,
length(profile.preferences) as pref_count,
profile.location.city as city
FROM Python(df_with_json)
SETTINGS pandas_analyze_sample = 1000 -- Analyze all rows for JSON detection
""", "Pretty")
print(result)
# Advanced JSON operations
complex_json = chdb.query("""
SELECT
user_id,
JSONLength(toString(profile)) as json_fields,
JSONType(toString(profile), 'preferences') as pref_type,
if(
JSONHas(toString(profile), 'achievements'),
JSONExtractString(toString(profile), 'achievements[0].title'),
'None'
) as first_achievement
FROM Python(df_with_json)
""", "JSONEachRow")
print(complex_json)