这是indexloc提供的服务，不要输入任何密码
Skip to content

Support non-computed nested array globals. #8611

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 2 commits into
base: master
Choose a base branch
from
Draft

Support non-computed nested array globals. #8611

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c152b45
Enable array of array globals.
dnwpark Apr 21, 2025
32db320
Support array of array globals.
dnwpark Mar 12, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 8 additions & 2 deletions edb/edgeql/compiler/expr.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1047,7 +1047,9 @@ def _infer_slice_type(
def compile_Indirection(
expr: qlast.Indirection, *, ctx: context.ContextLevel
) -> irast.Set:
node: irast.Set | irast.Expr = dispatch.compile(expr.arg, ctx=ctx)
node: irast.Set | irast.IndexIndirection | irast.SliceIndirection = (
dispatch.compile(expr.arg, ctx=ctx)
)
for indirection_el in expr.indirection:
if isinstance(indirection_el, qlast.Index):
idx = dispatch.compile(indirection_el.index, ctx=ctx)
Expand All @@ -1057,7 +1059,11 @@ def compile_Indirection(
)

node = irast.IndexIndirection(
expr=node, index=idx, typeref=typeref, span=expr.span
expr=node,
index=idx,
input_typeref=node.typeref,
typeref=typeref,
span=expr.span,
)
elif isinstance(indirection_el, qlast.Slice):
start: Optional[irast.Base]
Expand Down
37 changes: 36 additions & 1 deletion edb/edgeql/compiler/setgen.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2043,7 +2043,41 @@ def get_global_param(
param_name = f'__edb_global_{len(ctx.env.query_globals)}__'

target = glob.get_target(ctx.env.schema)
target_typeref = typegen.type_to_typeref(target, env=ctx.env)

target_typeref: irast.TypeRef
if (
glob.get_expr(ctx.env.schema) is None
and target.contains_array_of_array(ctx.env.schema)
):
# User specified global with array<array<...>> should treat it
# as array<tuple<array<...>>>
padded_collection_type: s_types.Collection
if isinstance(target, s_types.Array):
padded_collection_type = s_types.Array.get_padded_collection(
ctx.env.schema,
element_type=target.get_element_type(ctx.env.schema),
)
elif isinstance(target, s_types.Tuple):
padded_collection_type = s_types.Tuple.get_padded_collection(
ctx.env.schema,
element_types={
n: st
for n, st in (
target
.get_element_types(ctx.env.schema)
.items(ctx.env.schema)
)
},
named=target.get_named(ctx.env.schema),
)
else:
raise NotImplementedError
target_typeref = typegen.type_to_typeref(
padded_collection_type, env=ctx.env
)

else:
target_typeref = typegen.type_to_typeref(target, env=ctx.env)

ctx.env.query_globals[name] = irast.Global(
name=param_name,
Expand Down Expand Up @@ -2073,6 +2107,7 @@ def get_global_param_sets(
typeref=param.ir_type,
is_implicit_global=is_implicit_global,
),
type_override=param.schema_type,
ctx=ctx,
)
if glob.needs_present_arg(ctx.env.schema):
Expand Down
8 changes: 7 additions & 1 deletion edb/ir/ast.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -160,6 +160,11 @@ class TypeRef(ImmutableBase):
collection: typing.Optional[str] = None
# Collection subtypes if this is a collection
subtypes: tuple[TypeRef, ...] = ()
# If this is array<array<...>>, the padded type is array<tuple<array<...>>>
padded_array_type: typing.Optional[TypeRef] = None
# If this is array<...> and tuple<array<...>> is persisted
# in the schema, this is the tuple's id
wrapped_array_id: typing.Optional[uuid.UUID] = None
# True, if this describes a scalar type
is_scalar: bool = False
# True, if this describes a view
Expand Down Expand Up @@ -1081,8 +1086,9 @@ class OperatorCall(Call):

class IndexIndirection(ImmutableExpr):

expr: Base
expr: Set | IndexIndirection | SliceIndirection
index: Base
input_typeref: TypeRef
typeref: TypeRef


Expand Down
46 changes: 45 additions & 1 deletion edb/ir/typeutils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,6 +116,11 @@ def is_array(typeref: irast.TypeRef) -> bool:
return typeref.collection == s_types.Array.get_schema_name()


def is_array_of_array(typeref: irast.TypeRef) -> bool:
"""Return True if *typeref* describes an array of array type."""
return is_array(typeref) and is_array(typeref.subtypes[0])


def is_tuple(typeref: irast.TypeRef) -> bool:
"""Return True if *typeref* describes a tuple type."""
return typeref.collection == s_types.Tuple.get_schema_name()
Expand Down Expand Up @@ -191,6 +196,18 @@ def is_persistent_tuple(typeref: irast.TypeRef) -> bool:
return False


def is_persistent_array_of_array(typeref: irast.TypeRef) -> bool:
if is_array_of_array(typeref):
if typeref.material_type is not None:
material = typeref.material_type
else:
material = typeref

return material.in_schema
else:
return False


def is_empty_typeref(typeref: irast.TypeRef) -> bool:
return typeref.union is not None and len(typeref.union) == 0

Expand Down Expand Up @@ -480,6 +497,31 @@ def _typeref(
else:
material_typeref = None

padded_array_typeref = None
if (
isinstance(t, s_types.Array)
and (element_type := t.get_element_type(schema))
and isinstance(element_type, s_types.Array)
):
padded_collection_type = s_types.Array.get_padded_collection(
schema, element_type=element_type
)
padded_array_typeref = _typeref(padded_collection_type)

wrapped_array_id = None
if isinstance(t, s_types.Array):
wrapped_array_name = s_types.Tuple.generate_name(
element_names={'f1': t.get_name(schema)},
named=True,
)
if (
(wrapped_array_type := schema.get_global(
s_types.Tuple, wrapped_array_name, None
))
and wrapped_array_type.get_is_persistent(schema)
):
wrapped_array_id = wrapped_array_type.id

result = irast.TypeRef(
id=t.id,
name_hint=name_hint,
Expand All @@ -490,7 +532,9 @@ def _typeref(
in_schema=t.get_is_persistent(schema),
subtypes=tuple(
_typeref(st) for st in t.get_subtypes(schema)
)
),
padded_array_type=padded_array_typeref,
wrapped_array_id=wrapped_array_id,
)

if cache is not None and typename is None and _name is None:
Expand Down
179 changes: 104 additions & 75 deletions edb/pgsql/compiler/astutils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -139,84 +139,113 @@ def array_get_inner_array(
an element is accessed, it needs to be unwrapped.

Essentially, this function takes tuple<array<...>> and returns array<...>

Postgres does not support arbitrarily accessing fields out of unnamed
composites and so we need to do an extra unnest(array[]) to be able to
specify the name and type our resulting array.

For example, the query: `select [[1]][0];` will produce the following SQL:

SELECT
"expr-6~2"."array_value~4" AS "array_serialized~1"
FROM
LATERAL
(SELECT
"expr-5~2"."array_value~3" AS "array_value~4"
FROM
LATERAL
(SELECT
(SELECT
"0"
FROM
-- EXTRA unnest(array[])
unnest(ARRAY[
-- INDEX INDIRECTION
edgedb_v7_2f26206480._index(
"expr-3~2"."array_value~2",
($2)::int8,
'ERROR MESSAGE'
)
]) AS ("0" int8[])
) AS "array_value~3"
FROM
LATERAL
-- INITAL ARRAY [[1]]
(SELECT
ARRAY[ROW("expr-2~2"."array_value~1")]
AS "array_value~2"
FROM
LATERAL
(SELECT
ARRAY[($1)::int8]
AS "array_value~1"
) AS "expr-2~2"
) AS "expr-3~2"
) AS "expr-5~2"
) AS "expr-6~2"
WHERE
("expr-6~2"."array_value~4" IS NOT NULL)
LIMIT
(SELECT
(101)::int8 AS "expr~7_value~1"
)
"""
return pgast.SelectStmt(
target_list=[
pgast.ResTarget(val=pgast.ColumnRef(name=['0'])),
],
from_clause=[
pgast.RangeFunction(
functions=[
pgast.FuncCall(
name=('unnest',),
args=[
pgast.ArrayExpr(
elements=[wrapped_array],
)
],
coldeflist=[
pgast.ColumnDef(
name='0',
typename=pgast.TypeName(
name=pg_types.pg_type_from_ir_typeref(array_typeref)
# Nested arrays wrap their inner arrays in a tuple
if array_typeref.wrapped_array_id:
# temporary array of array implemented as array<tuple<array<...>>>
# using named tuple
return pgast.SelectStmt(
target_list=[
pgast.ResTarget(val=pgast.ColumnRef(name=['f1'])),
],
from_clause=[
pgast.RangeFunction(
functions=[
pgast.FuncCall(
name=('unnest',),
args=[
pgast.ArrayExpr(
elements=[wrapped_array],
)
)
]
)
]
],
)
]
)
]
)
else:
# temporary array of array implemented as array<tuple<array<...>>>
# using unnamed tuple
"""
Postgres does not support arbitrarily accessing fields out of unnamed
composites and so we need to do an extra unnest(array[]) to be able to
specify the name and type our resulting array.

For example, the query: `select [[1]][0];` will produce the following
SQL:

SELECT
"expr-6~2"."array_value~4" AS "array_serialized~1"
FROM
LATERAL
(SELECT
"expr-5~2"."array_value~3" AS "array_value~4"
FROM
LATERAL
(SELECT
(SELECT
"0"
FROM
-- EXTRA unnest(array[])
unnest(ARRAY[
-- INDEX INDIRECTION
edgedb_v7_2f26206480._index(
"expr-3~2"."array_value~2",
($2)::int8,
'ERROR MESSAGE'
)
]) AS ("0" int8[])
) AS "array_value~3"
FROM
LATERAL
-- INITAL ARRAY [[1]]
(SELECT
ARRAY[ROW("expr-2~2"."array_value~1")]
AS "array_value~2"
FROM
LATERAL
(SELECT
ARRAY[($1)::int8]
AS "array_value~1"
) AS "expr-2~2"
) AS "expr-3~2"
) AS "expr-5~2"
Comment on lines +181 to +212
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I know this is from a previous PR, but trim the comment example code down. It includes some extra irrelevant stuff: an extra wrapper to produce serialized, the toplevel NULL filter, an inject LIMIT

) AS "expr-6~2"
WHERE
("expr-6~2"."array_value~4" IS NOT NULL)
LIMIT
(SELECT
(101)::int8 AS "expr~7_value~1"
)
]
)
"""

return pgast.SelectStmt(
target_list=[
pgast.ResTarget(val=pgast.ColumnRef(name=['0'])),
],
from_clause=[
pgast.RangeFunction(
functions=[
pgast.FuncCall(
name=('unnest',),
args=[
pgast.ArrayExpr(
elements=[wrapped_array],
)
],
coldeflist=[
pgast.ColumnDef(
name='0',
typename=pgast.TypeName(
name=pg_types.pg_type_from_ir_typeref(array_typeref)
)
)
]
)
]
)
]
)


def is_null_const(expr: pgast.BaseExpr) -> bool:
Expand Down
22 changes: 14 additions & 8 deletions edb/pgsql/compiler/output.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -248,14 +248,20 @@ def array_as_json_object(
# To prevent this, we need to explicitly serialize the inner arrays then
# aggregate them.
el_name = 'f1'
coldeflist = [
pgast.ColumnDef(
name=str(el_name),
typename=pgast.TypeName(
name=pgtypes.pg_type_from_ir_typeref(el_type),
),
)
]
if (
styperef.padded_array_type is not None
and styperef.padded_array_type.in_schema
):
coldeflist = None
else:
coldeflist = [
pgast.ColumnDef(
name=str(el_name),
typename=pgast.TypeName(
name=pgtypes.pg_type_from_ir_typeref(el_type),
),
)
]
unwrapped_inner_array = pgast.RangeFunction(
functions=[
pgast.FuncCall(
Expand Down
Loading
Loading