Skip to content

Commit 7dd9702

Browse files
rustyconoverrustyconover
committed
fixes
1 parent 8ec6367 commit 7dd9702

2 files changed

Lines changed: 266 additions & 0 deletions

File tree

‎src/query_farm_server_base/duckdb_serialized_values.py‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,14 @@ def sql(self) -> str:
141141
return self.id
142142

143143

144+
class SerializedValueType_any(BaseModel):
145+
id: Literal["ANY"] = "ANY"
146+
type_info: None
147+
148+
def sql(self) -> str:
149+
raise NotImplementedError("SerializedValueType_any.sql() is not implemented.")
150+
151+
144152
class SerializedValue_boolean(SerializedValueBase):
145153
type: SerializedValueType_boolean
146154
value: bool | None = None
@@ -808,6 +816,7 @@ def get_discriminator_value(v: Any) -> str:
808816

809817

810818
AllValidTypeIdAndInfo = Union[
819+
SerializedValueType_any,
811820
SerializedValueType_boolean,
812821
SerializedValueType_bigint,
813822
SerializedValueType_bit,
@@ -820,6 +829,7 @@ def get_discriminator_value(v: Any) -> str:
820829
SerializedValueType_integer,
821830
SerializedValueType_interval,
822831
SerializedValueType_list,
832+
"SerializedValueType_map",
823833
SerializedValueType_null,
824834
SerializedValueType_smallint,
825835
"SerializedValueType_struct",
@@ -888,6 +898,50 @@ def sql(self) -> str:
888898
)
889899

890900

901+
class SerializedValueTypeInfo_map(BaseModel):
902+
child_type: SerializedValueType_struct
903+
type: Literal["LIST_TYPE_INFO"] = "LIST_TYPE_INFO"
904+
alias: str | None = None
905+
modifiers: list[Any] | None = None
906+
907+
908+
class SerializedValueType_map(BaseModel):
909+
id: Literal["MAP"] = "MAP"
910+
type_info: SerializedValueTypeInfo_map
911+
912+
def sql(self) -> str:
913+
return (
914+
"MAP("
915+
+ ",".join(
916+
[
917+
f'"{child.first}" {child.second.sql()}'
918+
for child in self.type_info.child_type.type_info.child_types
919+
]
920+
)
921+
+ ")"
922+
)
923+
924+
925+
class SerializedValueValue_map(BaseModel):
926+
children: list["SerializedValue"]
927+
928+
929+
class SerializedValue_map(SerializedValueBase):
930+
type: SerializedValueType_map
931+
value: SerializedValueValue_map
932+
933+
def sql(self) -> str:
934+
names = [child.first for child in self.type.type_info.child_type.type_info.child_types]
935+
values = self.value.children
936+
return (
937+
"{"
938+
+ ",".join(
939+
[f"'{name}':" + value.sql() for name, value in zip(names, values, strict=True)]
940+
)
941+
+ "}"
942+
)
943+
944+
891945
SerializedValue = Annotated[
892946
Annotated[SerializedValue_bigint, Tag("BIGINT")]
893947
| Annotated[SerializedValue_bit, Tag("BIT")]
@@ -901,6 +955,7 @@ def sql(self) -> str:
901955
| Annotated[SerializedValue_integer, Tag("INTEGER")]
902956
| Annotated[SerializedValue_interval, Tag("INTERVAL")]
903957
| Annotated[SerializedValue_list, Tag("LIST")]
958+
| Annotated[SerializedValue_list, Tag("MAP")]
904959
| Annotated[SerializedValue_null, Tag("NULL")]
905960
| Annotated[SerializedValue_smallint, Tag("SMALLINT")]
906961
| Annotated[SerializedValue_struct, Tag("STRUCT")]

‎src/query_farm_server_base/test_duckdb_serialized_expression.py‎

Lines changed: 211 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,217 @@
2020
@pytest.mark.parametrize(
2121
"parser_cls,input",
2222
[
23+
(
24+
duckdb_serialized_values.SerializedValueType_map,
25+
{
26+
"id": "MAP",
27+
"type_info": {
28+
"type": "LIST_TYPE_INFO",
29+
"alias": "",
30+
"extension_info": None,
31+
"child_type": {
32+
"id": "STRUCT",
33+
"type_info": {
34+
"type": "STRUCT_TYPE_INFO",
35+
"alias": "",
36+
"extension_info": None,
37+
"child_types": [
38+
{"first": "key", "second": {"id": "VARCHAR", "type_info": None}},
39+
{"first": "value", "second": {"id": "VARCHAR", "type_info": None}},
40+
],
41+
},
42+
},
43+
},
44+
},
45+
),
46+
(
47+
duckdb_serialized_values.ExpressionBoundColumnRef,
48+
{
49+
"expression_class": "BOUND_COLUMN_REF",
50+
"type": "BOUND_COLUMN_REF",
51+
"alias": "mapping",
52+
"query_location": 18446744073709551615,
53+
"return_type": {
54+
"id": "MAP",
55+
"type_info": {
56+
"type": "LIST_TYPE_INFO",
57+
"alias": "",
58+
"extension_info": None,
59+
"child_type": {
60+
"id": "STRUCT",
61+
"type_info": {
62+
"type": "STRUCT_TYPE_INFO",
63+
"alias": "",
64+
"extension_info": None,
65+
"child_types": [
66+
{
67+
"first": "key",
68+
"second": {"id": "VARCHAR", "type_info": None},
69+
},
70+
{
71+
"first": "value",
72+
"second": {"id": "VARCHAR", "type_info": None},
73+
},
74+
],
75+
},
76+
},
77+
},
78+
},
79+
"binding": {"table_index": 0, "column_index": 0},
80+
"depth": 0,
81+
},
82+
),
83+
(
84+
duckdb_serialized_values.ExpressionBoundFunction,
85+
{
86+
"expression_class": "BOUND_FUNCTION",
87+
"type": "BOUND_FUNCTION",
88+
"alias": "",
89+
"query_location": 39,
90+
"return_type": {"id": "VARCHAR", "type_info": None},
91+
"children": [
92+
{
93+
"expression_class": "BOUND_COLUMN_REF",
94+
"type": "BOUND_COLUMN_REF",
95+
"alias": "mapping",
96+
"query_location": 18446744073709551615,
97+
"return_type": {
98+
"id": "MAP",
99+
"type_info": {
100+
"type": "LIST_TYPE_INFO",
101+
"alias": "",
102+
"extension_info": None,
103+
"child_type": {
104+
"id": "STRUCT",
105+
"type_info": {
106+
"type": "STRUCT_TYPE_INFO",
107+
"alias": "",
108+
"extension_info": None,
109+
"child_types": [
110+
{
111+
"first": "key",
112+
"second": {"id": "VARCHAR", "type_info": None},
113+
},
114+
{
115+
"first": "value",
116+
"second": {"id": "VARCHAR", "type_info": None},
117+
},
118+
],
119+
},
120+
},
121+
},
122+
},
123+
"binding": {"table_index": 0, "column_index": 0},
124+
"depth": 0,
125+
},
126+
{
127+
"expression_class": "BOUND_CONSTANT",
128+
"type": "VALUE_CONSTANT",
129+
"alias": "",
130+
"query_location": 18446744073709551615,
131+
"value": {
132+
"type": {"id": "VARCHAR", "type_info": None},
133+
"is_null": False,
134+
"value": "color",
135+
},
136+
},
137+
],
138+
"name": "map_extract_value",
139+
"arguments": [
140+
{"id": "ANY", "type_info": None},
141+
{"id": "VARCHAR", "type_info": None},
142+
],
143+
"original_arguments": [],
144+
"catalog_name": "system",
145+
"schema_name": "main",
146+
"has_serialize": False,
147+
"is_operator": False,
148+
},
149+
),
150+
(
151+
duckdb_serialized_values.ExpressionBoundComparison,
152+
{
153+
"expression_class": "BOUND_COMPARISON",
154+
"type": "COMPARE_EQUAL",
155+
"alias": "",
156+
"query_location": 18446744073709551615,
157+
"left": {
158+
"expression_class": "BOUND_FUNCTION",
159+
"type": "BOUND_FUNCTION",
160+
"alias": "",
161+
"query_location": 39,
162+
"return_type": {"id": "VARCHAR", "type_info": None},
163+
"children": [
164+
{
165+
"expression_class": "BOUND_COLUMN_REF",
166+
"type": "BOUND_COLUMN_REF",
167+
"alias": "mapping",
168+
"query_location": 18446744073709551615,
169+
"return_type": {
170+
"id": "MAP",
171+
"type_info": {
172+
"type": "LIST_TYPE_INFO",
173+
"alias": "",
174+
"extension_info": None,
175+
"child_type": {
176+
"id": "STRUCT",
177+
"type_info": {
178+
"type": "STRUCT_TYPE_INFO",
179+
"alias": "",
180+
"extension_info": None,
181+
"child_types": [
182+
{
183+
"first": "key",
184+
"second": {"id": "VARCHAR", "type_info": None},
185+
},
186+
{
187+
"first": "value",
188+
"second": {"id": "VARCHAR", "type_info": None},
189+
},
190+
],
191+
},
192+
},
193+
},
194+
},
195+
"binding": {"table_index": 0, "column_index": 0},
196+
"depth": 0,
197+
},
198+
{
199+
"expression_class": "BOUND_CONSTANT",
200+
"type": "VALUE_CONSTANT",
201+
"alias": "",
202+
"query_location": 18446744073709551615,
203+
"value": {
204+
"type": {"id": "VARCHAR", "type_info": None},
205+
"is_null": False,
206+
"value": "color",
207+
},
208+
},
209+
],
210+
"name": "map_extract_value",
211+
"arguments": [
212+
{"id": "ANY", "type_info": None},
213+
{"id": "VARCHAR", "type_info": None},
214+
],
215+
"original_arguments": [],
216+
"catalog_name": "system",
217+
"schema_name": "main",
218+
"has_serialize": False,
219+
"is_operator": False,
220+
},
221+
"right": {
222+
"expression_class": "BOUND_CONSTANT",
223+
"type": "VALUE_CONSTANT",
224+
"alias": "",
225+
"query_location": 18446744073709551615,
226+
"value": {
227+
"type": {"id": "VARCHAR", "type_info": None},
228+
"is_null": False,
229+
"value": "red",
230+
},
231+
},
232+
},
233+
),
23234
(
24235
duckdb_serialized_values.SerializedValueType_list,
25236
{

0 commit comments

Comments
 (0)