You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "uranusjr (via GitHub)" <gi...@apache.org> on 2024/05/06 08:35:26 UTC
[PR] Unify lazy db sequence implementations [airflow]
uranusjr opened a new pull request, #39426:
URL: https://github.com/apache/airflow/pull/39426
LazyXComAccessor and InletEventsAccessor are somewhat similar in purpose. A new LazySelectSequence is introduced with a more generic implementation that fits both use cases. This also implements the Sequence protocol a bit more efficiently in some places.
See #39367 for some context.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Unify lazy db sequence implementations [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #39426:
URL: https://github.com/apache/airflow/pull/39426#discussion_r1592965060
##########
airflow/utils/db.py:
##########
@@ -1888,12 +1914,119 @@ def get_query_count(query_stmt: Select, *, session: Session) -> int:
return session.scalar(count_stmt)
+def get_query_exists(query_stmt: Select, *, session: Session) -> bool:
+ """Check whether there is at least one row matching a query.
+
+ A SELECT 1 FROM is issued against the subquery built from the given
+ statement. The ORDER BY clause is stripped from the statement since it's
+ unnecessary, and can impact query planning and degrade performance.
+
+ :meta private:
+ """
+ count_stmt = select(literal(True)).select_from(query_stmt.order_by(None).subquery())
+ return session.scalar(count_stmt)
+
+
def exists_query(*where: ClauseElement, session: Session) -> bool:
- """Check whether there is at least one row matching given clause.
+ """Check whether there is at least one row matching given clauses.
This does a SELECT 1 WHERE ... LIMIT 1 and check the result.
:meta private:
"""
stmt = select(literal(True)).where(*where).limit(1)
return session.scalar(stmt) is not None
+
+
+@attrs.define(slots=True)
+class LazySelectSequence(Sequence[T]):
+ """List-like interface to lazily access a database model query.
+
+ The intended use case is inside a task execution context, where we manage an
+ active SQLAlchemy session in the background.
+
+ :meta private:
+ """
+
+ _select_asc: ClauseElement
+ _select_desc: ClauseElement
+ _process_row: Callable[[Any], Any] = attrs.field(kw_only=True, default=operator.itemgetter(0))
+ _session: Session = attrs.field(kw_only=True, factory=get_current_task_instance_session)
+ _len: int | None = attrs.field(init=False, default=None)
+
+ @classmethod
+ def from_select(
+ cls,
+ select: Select,
+ *,
+ order_by: Sequence[ClauseElement],
+ process_row: Callable[[Any], Any] = operator.itemgetter(0),
+ session: Session | None = None,
+ ) -> Self:
+ s1 = select
+ for col in order_by:
+ s1 = s1.order_by(col.asc())
+ s2 = select
+ for col in order_by:
+ s2 = s2.order_by(col.desc())
+ return cls(s1, s2, process_row=process_row, session=session or get_current_task_instance_session())
+
+ def __repr__(self) -> str:
+ return f"LazySelectSequence([{len(self)} items])"
+
+ def __str__(self) -> str:
+ return str(list(self))
+
+ def __getstate__(self) -> Any:
+ # We don't want to go to the trouble of serializing SQLAlchemy objects.
+ # Converting the statement into a SQL string is the best we can get.
+ # The literal_binds compile argument inlines all the values into the SQL
+ # string to simplify cross-process commuinication as much as possible.
+ # Theoratically we can do the same for count(), but I think it should be
+ # performant enough to calculate only that eagerly.
+ s1 = str(self._select_asc.compile(self._session.get_bind(), compile_kwargs={"literal_binds": True}))
+ s2 = str(self._select_desc.compile(self._session.get_bind(), compile_kwargs={"literal_binds": True}))
+ return (s1, s2, self._process_row, len(self))
+
+ def __setstate__(self, state: Any) -> None:
+ s1, s2, self._process_row, self._len = state
+ self._select_asc = text(s1)
+ self._select_desc = text(s2)
+ self._session = get_current_task_instance_session()
+
+ def __bool__(self) -> bool:
+ return get_query_exists(self._select_asc, session=self._session)
+
+ def __eq__(self, other: Any) -> bool:
+ if not isinstance(other, collections.abc.Sequence):
+ return NotImplemented
+ z = itertools.zip_longest(iter(self), iter(other), fillvalue=object())
+ return all(x == y for x, y in z)
+
+ def __reversed__(self) -> Iterator[T]:
+ return iter(self._process_row(r) for r in self._session.execute(self._select_desc))
+
+ def __iter__(self) -> Iterator[T]:
+ return iter(self._process_row(r) for r in self._session.execute(self._select_asc))
+
+ def __len__(self) -> int:
+ if (le := self._len) is None:
+ le = self._len = get_query_count(self._select_asc, session=self._session)
+ return le
+
Review Comment:
I’ve removed the local variable.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Unify lazy db sequence implementations [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #39426:
URL: https://github.com/apache/airflow/pull/39426#discussion_r1593758556
##########
airflow/utils/db.py:
##########
@@ -1888,12 +1914,144 @@ def get_query_count(query_stmt: Select, *, session: Session) -> int:
return session.scalar(count_stmt)
+def get_query_exists(query_stmt: Select, *, session: Session) -> bool:
+ """Check whether there is at least one row matching a query.
+
+ A SELECT 1 FROM is issued against the subquery built from the given
+ statement. The ORDER BY clause is stripped from the statement since it's
+ unnecessary, and can impact query planning and degrade performance.
+
+ :meta private:
+ """
+ count_stmt = select(literal(True)).select_from(query_stmt.order_by(None).subquery())
+ return session.scalar(count_stmt)
+
+
def exists_query(*where: ClauseElement, session: Session) -> bool:
- """Check whether there is at least one row matching given clause.
+ """Check whether there is at least one row matching given clauses.
This does a SELECT 1 WHERE ... LIMIT 1 and check the result.
:meta private:
"""
stmt = select(literal(True)).where(*where).limit(1)
return session.scalar(stmt) is not None
+
+
+@attrs.define(slots=True)
+class LazySelectSequence(Sequence[T]):
+ """List-like interface to lazily access a database model query.
+
+ The intended use case is inside a task execution context, where we manage an
+ active SQLAlchemy session in the background.
+
+ This is an abstract base class. Each use case should subclass, and implement
+ the following static methods:
+
+ * ``_rebuild_select`` is called when a lazy sequence is unpickled. Since it
+ is not easy to pickle SQLAlchemy constructs, this class serializes the
+ SELECT statements into plain text to storage. This method is called on
+ deserialization to convert the textual clause back into an ORM SELECT.
+ * ``_process_row`` is called when an item is accessed. The lazy sequence
+ uses ``session.execute()`` to fetch rows from the database, and this
+ method should know how to process each row into a value.
+
+ :meta private:
+ """
+
+ _select_asc: ClauseElement
+ _select_desc: ClauseElement
+ _session: Session = attrs.field(kw_only=True, factory=get_current_task_instance_session)
+ _len: int | None = attrs.field(init=False, default=None)
+
+ @classmethod
+ def from_select(
+ cls,
+ select: Select,
+ *,
+ order_by: Sequence[ClauseElement],
+ session: Session | None = None,
+ ) -> Self:
+ s1 = select
+ for col in order_by:
+ s1 = s1.order_by(col.asc())
+ s2 = select
+ for col in order_by:
+ s2 = s2.order_by(col.desc())
+ return cls(s1, s2, session=session or get_current_task_instance_session())
+
+ @staticmethod
+ def _rebuild_select(stmt: TextClause) -> Select:
+ """Rebuild a textual statement into an ORM-configured SELECT statement.
+
+ This should do something like ``select(field).from_statement(stmt)`` to
+ reconfigure ORM information to the textual SQL statement.
+ """
+ raise NotImplementedError
+
+ @staticmethod
+ def _process_row(row: Row) -> T:
+ """Process a SELECT-ed row into the end value."""
+ raise NotImplementedError
+
+ def __repr__(self) -> str:
+ counter = "item" if (length := len(self)) == 1 else "items"
+ return f"LazySelectSequence([{length} {counter}])"
+
+ def __str__(self) -> str:
+ counter = "item" if (length := len(self)) == 1 else "items"
+ return f"[... {length} {counter} ...]"
Review Comment:
It is not particularly useful to want to print the _entire_ sequence in a task, and when you do it (likely out of desperation or something) resolving the list would be extremely costly for most things. In most cases you’d only want the last few items.
It is not unreasonable for a DatasetEvent sequence to contain hundreds of rows, printing that all out is a disaster in more than one way, and I feel this might actually be better for debugging in that case.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Unify lazy db sequence implementations [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #39426:
URL: https://github.com/apache/airflow/pull/39426#discussion_r1593834483
##########
airflow/utils/db.py:
##########
@@ -1888,12 +1914,144 @@ def get_query_count(query_stmt: Select, *, session: Session) -> int:
return session.scalar(count_stmt)
+def get_query_exists(query_stmt: Select, *, session: Session) -> bool:
+ """Check whether there is at least one row matching a query.
+
+ A SELECT 1 FROM is issued against the subquery built from the given
+ statement. The ORDER BY clause is stripped from the statement since it's
+ unnecessary, and can impact query planning and degrade performance.
+
+ :meta private:
+ """
+ count_stmt = select(literal(True)).select_from(query_stmt.order_by(None).subquery())
+ return session.scalar(count_stmt)
+
+
def exists_query(*where: ClauseElement, session: Session) -> bool:
- """Check whether there is at least one row matching given clause.
+ """Check whether there is at least one row matching given clauses.
This does a SELECT 1 WHERE ... LIMIT 1 and check the result.
:meta private:
"""
stmt = select(literal(True)).where(*where).limit(1)
return session.scalar(stmt) is not None
+
+
+@attrs.define(slots=True)
+class LazySelectSequence(Sequence[T]):
+ """List-like interface to lazily access a database model query.
+
+ The intended use case is inside a task execution context, where we manage an
+ active SQLAlchemy session in the background.
+
+ This is an abstract base class. Each use case should subclass, and implement
+ the following static methods:
+
+ * ``_rebuild_select`` is called when a lazy sequence is unpickled. Since it
+ is not easy to pickle SQLAlchemy constructs, this class serializes the
+ SELECT statements into plain text to storage. This method is called on
+ deserialization to convert the textual clause back into an ORM SELECT.
+ * ``_process_row`` is called when an item is accessed. The lazy sequence
+ uses ``session.execute()`` to fetch rows from the database, and this
+ method should know how to process each row into a value.
+
+ :meta private:
+ """
+
+ _select_asc: ClauseElement
+ _select_desc: ClauseElement
+ _session: Session = attrs.field(kw_only=True, factory=get_current_task_instance_session)
+ _len: int | None = attrs.field(init=False, default=None)
+
+ @classmethod
+ def from_select(
+ cls,
+ select: Select,
+ *,
+ order_by: Sequence[ClauseElement],
+ session: Session | None = None,
+ ) -> Self:
+ s1 = select
+ for col in order_by:
+ s1 = s1.order_by(col.asc())
+ s2 = select
+ for col in order_by:
+ s2 = s2.order_by(col.desc())
+ return cls(s1, s2, session=session or get_current_task_instance_session())
+
+ @staticmethod
+ def _rebuild_select(stmt: TextClause) -> Select:
+ """Rebuild a textual statement into an ORM-configured SELECT statement.
+
+ This should do something like ``select(field).from_statement(stmt)`` to
+ reconfigure ORM information to the textual SQL statement.
+ """
+ raise NotImplementedError
+
+ @staticmethod
+ def _process_row(row: Row) -> T:
+ """Process a SELECT-ed row into the end value."""
+ raise NotImplementedError
+
+ def __repr__(self) -> str:
+ counter = "item" if (length := len(self)) == 1 else "items"
+ return f"LazySelectSequence([{length} {counter}])"
+
+ def __str__(self) -> str:
+ counter = "item" if (length := len(self)) == 1 else "items"
+ return f"[... {length} {counter} ...]"
Review Comment:
Added the class name back.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Unify lazy db sequence implementations [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr merged PR #39426:
URL: https://github.com/apache/airflow/pull/39426
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Unify lazy db sequence implementations [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #39426:
URL: https://github.com/apache/airflow/pull/39426#discussion_r1593765014
##########
airflow/utils/db.py:
##########
@@ -1888,12 +1914,144 @@ def get_query_count(query_stmt: Select, *, session: Session) -> int:
return session.scalar(count_stmt)
+def get_query_exists(query_stmt: Select, *, session: Session) -> bool:
+ """Check whether there is at least one row matching a query.
+
+ A SELECT 1 FROM is issued against the subquery built from the given
+ statement. The ORDER BY clause is stripped from the statement since it's
+ unnecessary, and can impact query planning and degrade performance.
+
+ :meta private:
+ """
+ count_stmt = select(literal(True)).select_from(query_stmt.order_by(None).subquery())
+ return session.scalar(count_stmt)
+
+
def exists_query(*where: ClauseElement, session: Session) -> bool:
- """Check whether there is at least one row matching given clause.
+ """Check whether there is at least one row matching given clauses.
This does a SELECT 1 WHERE ... LIMIT 1 and check the result.
:meta private:
"""
stmt = select(literal(True)).where(*where).limit(1)
return session.scalar(stmt) is not None
+
+
+@attrs.define(slots=True)
+class LazySelectSequence(Sequence[T]):
+ """List-like interface to lazily access a database model query.
+
+ The intended use case is inside a task execution context, where we manage an
+ active SQLAlchemy session in the background.
+
+ This is an abstract base class. Each use case should subclass, and implement
+ the following static methods:
+
+ * ``_rebuild_select`` is called when a lazy sequence is unpickled. Since it
+ is not easy to pickle SQLAlchemy constructs, this class serializes the
+ SELECT statements into plain text to storage. This method is called on
+ deserialization to convert the textual clause back into an ORM SELECT.
+ * ``_process_row`` is called when an item is accessed. The lazy sequence
+ uses ``session.execute()`` to fetch rows from the database, and this
+ method should know how to process each row into a value.
+
+ :meta private:
+ """
+
+ _select_asc: ClauseElement
+ _select_desc: ClauseElement
+ _session: Session = attrs.field(kw_only=True, factory=get_current_task_instance_session)
+ _len: int | None = attrs.field(init=False, default=None)
+
+ @classmethod
+ def from_select(
+ cls,
+ select: Select,
+ *,
+ order_by: Sequence[ClauseElement],
+ session: Session | None = None,
+ ) -> Self:
+ s1 = select
+ for col in order_by:
+ s1 = s1.order_by(col.asc())
+ s2 = select
+ for col in order_by:
+ s2 = s2.order_by(col.desc())
+ return cls(s1, s2, session=session or get_current_task_instance_session())
+
+ @staticmethod
+ def _rebuild_select(stmt: TextClause) -> Select:
+ """Rebuild a textual statement into an ORM-configured SELECT statement.
+
+ This should do something like ``select(field).from_statement(stmt)`` to
+ reconfigure ORM information to the textual SQL statement.
+ """
+ raise NotImplementedError
+
+ @staticmethod
+ def _process_row(row: Row) -> T:
+ """Process a SELECT-ed row into the end value."""
+ raise NotImplementedError
+
+ def __repr__(self) -> str:
+ counter = "item" if (length := len(self)) == 1 else "items"
+ return f"LazySelectSequence([{length} {counter}])"
+
+ def __str__(self) -> str:
+ counter = "item" if (length := len(self)) == 1 else "items"
+ return f"[... {length} {counter} ...]"
Review Comment:
Totally agree that we shouldn't return all the items, but I feel like including `LazySelectSequence` might be useful. What do you think?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Unify lazy db sequence implementations [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #39426:
URL: https://github.com/apache/airflow/pull/39426#discussion_r1593711216
##########
airflow/utils/db.py:
##########
@@ -1888,12 +1914,144 @@ def get_query_count(query_stmt: Select, *, session: Session) -> int:
return session.scalar(count_stmt)
+def get_query_exists(query_stmt: Select, *, session: Session) -> bool:
+ """Check whether there is at least one row matching a query.
+
+ A SELECT 1 FROM is issued against the subquery built from the given
+ statement. The ORDER BY clause is stripped from the statement since it's
+ unnecessary, and can impact query planning and degrade performance.
+
+ :meta private:
+ """
+ count_stmt = select(literal(True)).select_from(query_stmt.order_by(None).subquery())
+ return session.scalar(count_stmt)
+
+
def exists_query(*where: ClauseElement, session: Session) -> bool:
- """Check whether there is at least one row matching given clause.
+ """Check whether there is at least one row matching given clauses.
This does a SELECT 1 WHERE ... LIMIT 1 and check the result.
:meta private:
"""
stmt = select(literal(True)).where(*where).limit(1)
return session.scalar(stmt) is not None
+
+
+@attrs.define(slots=True)
+class LazySelectSequence(Sequence[T]):
+ """List-like interface to lazily access a database model query.
+
+ The intended use case is inside a task execution context, where we manage an
+ active SQLAlchemy session in the background.
+
+ This is an abstract base class. Each use case should subclass, and implement
+ the following static methods:
+
+ * ``_rebuild_select`` is called when a lazy sequence is unpickled. Since it
+ is not easy to pickle SQLAlchemy constructs, this class serializes the
+ SELECT statements into plain text to storage. This method is called on
+ deserialization to convert the textual clause back into an ORM SELECT.
+ * ``_process_row`` is called when an item is accessed. The lazy sequence
+ uses ``session.execute()`` to fetch rows from the database, and this
+ method should know how to process each row into a value.
+
+ :meta private:
+ """
+
+ _select_asc: ClauseElement
+ _select_desc: ClauseElement
+ _session: Session = attrs.field(kw_only=True, factory=get_current_task_instance_session)
+ _len: int | None = attrs.field(init=False, default=None)
+
+ @classmethod
+ def from_select(
+ cls,
+ select: Select,
+ *,
+ order_by: Sequence[ClauseElement],
+ session: Session | None = None,
+ ) -> Self:
+ s1 = select
+ for col in order_by:
+ s1 = s1.order_by(col.asc())
+ s2 = select
+ for col in order_by:
+ s2 = s2.order_by(col.desc())
+ return cls(s1, s2, session=session or get_current_task_instance_session())
+
+ @staticmethod
+ def _rebuild_select(stmt: TextClause) -> Select:
+ """Rebuild a textual statement into an ORM-configured SELECT statement.
+
+ This should do something like ``select(field).from_statement(stmt)`` to
+ reconfigure ORM information to the textual SQL statement.
+ """
+ raise NotImplementedError
+
+ @staticmethod
+ def _process_row(row: Row) -> T:
+ """Process a SELECT-ed row into the end value."""
+ raise NotImplementedError
+
+ def __repr__(self) -> str:
+ counter = "item" if (length := len(self)) == 1 else "items"
+ return f"LazySelectSequence([{length} {counter}])"
+
+ def __str__(self) -> str:
+ counter = "item" if (length := len(self)) == 1 else "items"
+ return f"[... {length} {counter} ...]"
Review Comment:
May I know why we are changing it this way? It looks like something hard to debug if we're trying to log or print it.
##########
airflow/utils/db.py:
##########
@@ -1888,12 +1914,144 @@ def get_query_count(query_stmt: Select, *, session: Session) -> int:
return session.scalar(count_stmt)
+def get_query_exists(query_stmt: Select, *, session: Session) -> bool:
Review Comment:
nitpick: I can't understand the meaning of this function at first glance. 👀 Not sure whether `check_query_exists` make more sense
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Unify lazy db sequence implementations [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #39426:
URL: https://github.com/apache/airflow/pull/39426#discussion_r1590842956
##########
airflow/models/base.py:
##########
@@ -48,7 +48,10 @@ def _get_schema():
mapper_registry = registry(metadata=metadata)
_sentinel = object()
-Base: Any = mapper_registry.generate_base()
+if TYPE_CHECKING:
+ Base = Any
+else:
+ Base = mapper_registry.generate_base()
Review Comment:
Not technically required, but needed to get past Mypy.
##########
airflow/models/xcom.py:
##########
@@ -688,16 +684,20 @@ def _deserialize_value(result: XCom, orm: bool) -> Any:
if orm:
object_hook = XComDecoder.orm_object_hook
- if result.value is None:
+ value = result.value
+ if isinstance(value, memoryview):
+ value = value.tobytes()
Review Comment:
A side effect of converting the XCom SQLAlchemy code to New Style. SQLAlchemy seems to return `memoryview` instead of `bytes` for binary blob columns (`XCom.value`).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Unify lazy db sequence implementations [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #39426:
URL: https://github.com/apache/airflow/pull/39426#discussion_r1593759656
##########
airflow/utils/db.py:
##########
@@ -1888,12 +1914,144 @@ def get_query_count(query_stmt: Select, *, session: Session) -> int:
return session.scalar(count_stmt)
+def get_query_exists(query_stmt: Select, *, session: Session) -> bool:
Review Comment:
I was following the naming of `get_query_count` (admittedly not a very good one either). Using `check` instead makes sense.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Unify lazy db sequence implementations [airflow]
Posted by "Lee-W (via GitHub)" <gi...@apache.org>.
Lee-W commented on code in PR #39426:
URL: https://github.com/apache/airflow/pull/39426#discussion_r1592111536
##########
airflow/utils/db.py:
##########
@@ -1888,12 +1914,119 @@ def get_query_count(query_stmt: Select, *, session: Session) -> int:
return session.scalar(count_stmt)
+def get_query_exists(query_stmt: Select, *, session: Session) -> bool:
+ """Check whether there is at least one row matching a query.
+
+ A SELECT 1 FROM is issued against the subquery built from the given
+ statement. The ORDER BY clause is stripped from the statement since it's
+ unnecessary, and can impact query planning and degrade performance.
+
+ :meta private:
+ """
+ count_stmt = select(literal(True)).select_from(query_stmt.order_by(None).subquery())
+ return session.scalar(count_stmt)
+
+
def exists_query(*where: ClauseElement, session: Session) -> bool:
- """Check whether there is at least one row matching given clause.
+ """Check whether there is at least one row matching given clauses.
This does a SELECT 1 WHERE ... LIMIT 1 and check the result.
:meta private:
"""
stmt = select(literal(True)).where(*where).limit(1)
return session.scalar(stmt) is not None
+
+
+@attrs.define(slots=True)
+class LazySelectSequence(Sequence[T]):
+ """List-like interface to lazily access a database model query.
+
+ The intended use case is inside a task execution context, where we manage an
+ active SQLAlchemy session in the background.
+
+ :meta private:
+ """
+
+ _select_asc: ClauseElement
+ _select_desc: ClauseElement
+ _process_row: Callable[[Any], Any] = attrs.field(kw_only=True, default=operator.itemgetter(0))
+ _session: Session = attrs.field(kw_only=True, factory=get_current_task_instance_session)
+ _len: int | None = attrs.field(init=False, default=None)
+
+ @classmethod
+ def from_select(
+ cls,
+ select: Select,
+ *,
+ order_by: Sequence[ClauseElement],
+ process_row: Callable[[Any], Any] = operator.itemgetter(0),
+ session: Session | None = None,
+ ) -> Self:
+ s1 = select
+ for col in order_by:
+ s1 = s1.order_by(col.asc())
+ s2 = select
+ for col in order_by:
+ s2 = s2.order_by(col.desc())
+ return cls(s1, s2, process_row=process_row, session=session or get_current_task_instance_session())
+
+ def __repr__(self) -> str:
+ return f"LazySelectSequence([{len(self)} items])"
+
+ def __str__(self) -> str:
+ return str(list(self))
+
+ def __getstate__(self) -> Any:
+ # We don't want to go to the trouble of serializing SQLAlchemy objects.
+ # Converting the statement into a SQL string is the best we can get.
+ # The literal_binds compile argument inlines all the values into the SQL
+ # string to simplify cross-process commuinication as much as possible.
+ # Theoratically we can do the same for count(), but I think it should be
+ # performant enough to calculate only that eagerly.
+ s1 = str(self._select_asc.compile(self._session.get_bind(), compile_kwargs={"literal_binds": True}))
+ s2 = str(self._select_desc.compile(self._session.get_bind(), compile_kwargs={"literal_binds": True}))
+ return (s1, s2, self._process_row, len(self))
+
+ def __setstate__(self, state: Any) -> None:
+ s1, s2, self._process_row, self._len = state
+ self._select_asc = text(s1)
+ self._select_desc = text(s2)
+ self._session = get_current_task_instance_session()
+
+ def __bool__(self) -> bool:
+ return get_query_exists(self._select_asc, session=self._session)
+
+ def __eq__(self, other: Any) -> bool:
+ if not isinstance(other, collections.abc.Sequence):
+ return NotImplemented
+ z = itertools.zip_longest(iter(self), iter(other), fillvalue=object())
+ return all(x == y for x, y in z)
+
+ def __reversed__(self) -> Iterator[T]:
+ return iter(self._process_row(r) for r in self._session.execute(self._select_desc))
+
+ def __iter__(self) -> Iterator[T]:
+ return iter(self._process_row(r) for r in self._session.execute(self._select_asc))
+
+ def __len__(self) -> int:
+ if (le := self._len) is None:
+ le = self._len = get_query_count(self._select_asc, session=self._session)
+ return le
+
Review Comment:
Is `le` required here? Or could we just return `self._len` instead
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Unify lazy db sequence implementations [airflow]
Posted by "ephraimbuddy (via GitHub)" <gi...@apache.org>.
ephraimbuddy commented on code in PR #39426:
URL: https://github.com/apache/airflow/pull/39426#discussion_r1593741775
##########
docs/apache-airflow/authoring-and-scheduling/dynamic-task-mapping.rst:
##########
@@ -53,7 +53,7 @@ The grid view also provides visibility into your mapped tasks in the details pan
In the above example, ``values`` received by ``sum_it`` is an aggregation of all values returned by each mapped instance of ``add_one``. However, since it is impossible to know how many instances of ``add_one`` we will have in advance, ``values`` is not a normal list, but a "lazy sequence" that retrieves each individual value only when asked. Therefore, if you run ``print(values)`` directly, you would get something like this::
- LazyXComAccess(dag_id='simple_mapping', run_id='test_run', task_id='add_one')
+ [... 15 items ...]
Review Comment:
Curious what this means
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Unify lazy db sequence implementations [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #39426:
URL: https://github.com/apache/airflow/pull/39426#discussion_r1593762549
##########
docs/apache-airflow/authoring-and-scheduling/dynamic-task-mapping.rst:
##########
@@ -53,7 +53,7 @@ The grid view also provides visibility into your mapped tasks in the details pan
In the above example, ``values`` received by ``sum_it`` is an aggregation of all values returned by each mapped instance of ``add_one``. However, since it is impossible to know how many instances of ``add_one`` we will have in advance, ``values`` is not a normal list, but a "lazy sequence" that retrieves each individual value only when asked. Therefore, if you run ``print(values)`` directly, you would get something like this::
- LazyXComAccess(dag_id='simple_mapping', run_id='test_run', task_id='add_one')
+ [... 15 items ...]
Review Comment:
This is what will appear in the logs. Suggestions are welcome, but as mentioned in a thread above I don’t think it’s a good idea to print the entire list here.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Unify lazy db sequence implementations [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #39426:
URL: https://github.com/apache/airflow/pull/39426#discussion_r1592596293
##########
airflow/utils/db.py:
##########
@@ -1888,12 +1914,119 @@ def get_query_count(query_stmt: Select, *, session: Session) -> int:
return session.scalar(count_stmt)
+def get_query_exists(query_stmt: Select, *, session: Session) -> bool:
+ """Check whether there is at least one row matching a query.
+
+ A SELECT 1 FROM is issued against the subquery built from the given
+ statement. The ORDER BY clause is stripped from the statement since it's
+ unnecessary, and can impact query planning and degrade performance.
+
+ :meta private:
+ """
+ count_stmt = select(literal(True)).select_from(query_stmt.order_by(None).subquery())
+ return session.scalar(count_stmt)
+
+
def exists_query(*where: ClauseElement, session: Session) -> bool:
- """Check whether there is at least one row matching given clause.
+ """Check whether there is at least one row matching given clauses.
This does a SELECT 1 WHERE ... LIMIT 1 and check the result.
:meta private:
"""
stmt = select(literal(True)).where(*where).limit(1)
return session.scalar(stmt) is not None
+
+
+@attrs.define(slots=True)
+class LazySelectSequence(Sequence[T]):
+ """List-like interface to lazily access a database model query.
+
+ The intended use case is inside a task execution context, where we manage an
+ active SQLAlchemy session in the background.
+
+ :meta private:
+ """
+
+ _select_asc: ClauseElement
+ _select_desc: ClauseElement
+ _process_row: Callable[[Any], Any] = attrs.field(kw_only=True, default=operator.itemgetter(0))
+ _session: Session = attrs.field(kw_only=True, factory=get_current_task_instance_session)
+ _len: int | None = attrs.field(init=False, default=None)
+
+ @classmethod
+ def from_select(
+ cls,
+ select: Select,
+ *,
+ order_by: Sequence[ClauseElement],
+ process_row: Callable[[Any], Any] = operator.itemgetter(0),
+ session: Session | None = None,
+ ) -> Self:
+ s1 = select
+ for col in order_by:
+ s1 = s1.order_by(col.asc())
+ s2 = select
+ for col in order_by:
+ s2 = s2.order_by(col.desc())
+ return cls(s1, s2, process_row=process_row, session=session or get_current_task_instance_session())
+
+ def __repr__(self) -> str:
+ return f"LazySelectSequence([{len(self)} items])"
+
+ def __str__(self) -> str:
+ return str(list(self))
+
+ def __getstate__(self) -> Any:
+ # We don't want to go to the trouble of serializing SQLAlchemy objects.
+ # Converting the statement into a SQL string is the best we can get.
+ # The literal_binds compile argument inlines all the values into the SQL
+ # string to simplify cross-process commuinication as much as possible.
+ # Theoratically we can do the same for count(), but I think it should be
+ # performant enough to calculate only that eagerly.
+ s1 = str(self._select_asc.compile(self._session.get_bind(), compile_kwargs={"literal_binds": True}))
+ s2 = str(self._select_desc.compile(self._session.get_bind(), compile_kwargs={"literal_binds": True}))
+ return (s1, s2, self._process_row, len(self))
+
+ def __setstate__(self, state: Any) -> None:
+ s1, s2, self._process_row, self._len = state
+ self._select_asc = text(s1)
+ self._select_desc = text(s2)
+ self._session = get_current_task_instance_session()
+
+ def __bool__(self) -> bool:
+ return get_query_exists(self._select_asc, session=self._session)
+
+ def __eq__(self, other: Any) -> bool:
+ if not isinstance(other, collections.abc.Sequence):
+ return NotImplemented
+ z = itertools.zip_longest(iter(self), iter(other), fillvalue=object())
+ return all(x == y for x, y in z)
+
+ def __reversed__(self) -> Iterator[T]:
+ return iter(self._process_row(r) for r in self._session.execute(self._select_desc))
+
+ def __iter__(self) -> Iterator[T]:
+ return iter(self._process_row(r) for r in self._session.execute(self._select_asc))
+
+ def __len__(self) -> int:
+ if (le := self._len) is None:
+ le = self._len = get_query_count(self._select_asc, session=self._session)
+ return le
+
Review Comment:
Not required
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Unify lazy db sequence implementations [airflow]
Posted by "uranusjr (via GitHub)" <gi...@apache.org>.
uranusjr commented on code in PR #39426:
URL: https://github.com/apache/airflow/pull/39426#discussion_r1592964841
##########
airflow/models/xcom.py:
##########
@@ -688,16 +684,20 @@ def _deserialize_value(result: XCom, orm: bool) -> Any:
if orm:
object_hook = XComDecoder.orm_object_hook
- if result.value is None:
+ value = result.value
+ if isinstance(value, memoryview):
+ value = value.tobytes()
Review Comment:
Found another way to do this instead.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org