pankajkoti commented on code in PR #74017:
URL: https://github.com/apache/airflow/pull/74017#discussion_r4158020801
##########
providers/common/ai/src/airflow/providers/common/ai/utils/query_results.py:
##########
@@ -155,3 +188,136 @@ def build_query_result(
f"than paging through the result."
)
return _dumps(output)
+
+
+def _build_type_histogram(columns: Sequence[dict[str, str]]) -> dict[str, int]:
+ """
+ Count columns per type, capped to the most common
``_SCHEMA_TYPE_HISTOGRAM_TOP_K`` types.
+
+ Ordered by ``(-count, type)`` so the output is content-stable
(deterministic for byte
+ accounting and tests, not merely input-ordered). The long tail is folded
into a single
+ aggregate entry rather than listed, so parametrized types cannot inflate
the key count.
+
+ :param columns: ``{"name", "type"}`` dicts.
+ """
+ counts = Counter(col["type"] for col in columns)
+ ordered = sorted(counts.items(), key=lambda item: (-item[1], item[0]))
+ if len(ordered) <= _SCHEMA_TYPE_HISTOGRAM_TOP_K:
+ return dict(ordered)
+ histogram = dict(ordered[:_SCHEMA_TYPE_HISTOGRAM_TOP_K])
+ histogram[_SCHEMA_TYPE_HISTOGRAM_OTHER_KEY] = sum(
+ count for _, count in ordered[_SCHEMA_TYPE_HISTOGRAM_TOP_K:]
+ )
+ return histogram
+
+
+def build_schema_result(
+ columns: Sequence[dict[str, str]],
+ *,
+ max_columns: int,
+ max_result_bytes: int,
+ name_contains: str | None = None,
+) -> str:
+ """
+ Render a table's columns as a bounded JSON tool result.
+
+ Column names are the information an agent needs to write SQL, so unlike
query rows they cannot
+ simply be dropped: above ``max_columns`` (or the byte budget) the full
list is replaced by a
+ summary -- count, a type histogram, and a sample of columns -- that names
``name_contains`` as
+ the way to retrieve specific columns. ``columns`` is assumed to hold
distinct names (a table's
+ introspected columns are unique by construction).
+
+ :param columns: ``{"name", "type"}`` dicts in table order.
+ :param max_columns: Column count above which a summary replaces the full
list.
+ :param max_result_bytes: Budget for the serialized result.
+ :param name_contains: Case-insensitive substring. When given (and
non-empty), only matching
+ columns are considered and the value is echoed back so a filtered
subset is never mistaken
+ for the whole table.
+ """
+ name_contains = name_contains or None
+ total_columns = len(columns)
+ if name_contains is not None:
+ needle = name_contains.casefold()
+ selected = [col for col in columns if needle in col["name"].casefold()]
+ else:
+ selected = list(columns)
+
+ if name_contains is not None and not selected:
+ plural = "" if total_columns == 1 else "s"
+ return _dumps(
+ {
+ "columns": [],
+ "column_count": 0,
+ "name_contains": name_contains,
+ "hint": (
+ f"No columns match name_contains={name_contains!r}. Call
get_schema without "
+ f"name_contains to list all {total_columns}
column{plural}."
+ ),
+ }
+ )
+
+ full: dict[str, Any] = {"columns": selected, "column_count": len(selected)}
+ if name_contains is not None:
+ full["name_contains"] = name_contains
+ full["total_columns"] = total_columns
+ if len(selected) <= max_columns and _size(full) <= max_result_bytes:
+ return _dumps(full)
+
+ return _summarize_schema(
+ selected,
+ total_columns=total_columns,
+ max_columns=max_columns,
+ max_result_bytes=max_result_bytes,
+ name_contains=name_contains,
+ )
+
+
+def _summarize_schema(
+ selected: list[dict[str, str]],
+ *,
+ total_columns: int,
+ max_columns: int,
+ max_result_bytes: int,
+ name_contains: str | None,
+) -> str:
+ """Build the bounded summary returned when the full column list does not
fit."""
+ # Count is checked before bytes: it is the cheaper, more explainable bound
and the one the
+ # issue is about. A bytes-only truncation then means "count fits but
names/types are
+ # pathologically long", a distinct and rarer signal.
+ truncated_by = "max_columns" if len(selected) > max_columns else
"max_result_bytes"
+ hint = (
+ f"This filter matches {len(selected)} columns, more than can be
returned at once. "
+ "Use a more specific name_contains substring to narrow to the columns
you need."
+ if name_contains is not None
+ else (
+ f"This table has {len(selected)} columns, more than can be
returned at once. Call "
+ "get_schema again with a name_contains substring to return only
the matching columns."
+ )
+ )
+
+ output: dict[str, Any] = {
+ "column_count": len(selected),
+ "truncated": True,
+ "truncated_by": truncated_by,
+ "hint": hint,
+ }
+ if name_contains is not None:
+ output["name_contains"] = name_contains
+ output["total_columns"] = total_columns
+
+ # The core above is the guaranteed-useful payload. Add the histogram and a
sample of columns
+ # only while they fit, reserving the core first -- the same
contiguous-prefix accounting
+ # build_query_result uses for rows -- so a pathologically small budget
still returns the core.
+ output["type_histogram"] = _build_type_histogram(selected)
+ output["sample_columns"] = []
+ if _size(output) > max_result_bytes:
+ del output["type_histogram"]
+ return _dumps(output)
+ budget = max_result_bytes - _size(output)
+ for col in selected[:max_columns]:
Review Comment:
Changed `sample_columns` to a small fixed preview (20, independent of
`max_columns`) instead of the first `max_columns` columns, so the summary stays
well under the full list at any width and reads as a naming hint rather than a
prefix. On the 101-column case it is now 1374 bytes vs 5950 for the full list.
`max_columns` stays as the count trigger.
##########
providers/common/ai/src/airflow/providers/common/ai/utils/query_results.py:
##########
@@ -155,3 +188,136 @@ def build_query_result(
f"than paging through the result."
)
return _dumps(output)
+
+
+def _build_type_histogram(columns: Sequence[dict[str, str]]) -> dict[str, int]:
+ """
+ Count columns per type, capped to the most common
``_SCHEMA_TYPE_HISTOGRAM_TOP_K`` types.
+
+ Ordered by ``(-count, type)`` so the output is content-stable
(deterministic for byte
+ accounting and tests, not merely input-ordered). The long tail is folded
into a single
+ aggregate entry rather than listed, so parametrized types cannot inflate
the key count.
+
+ :param columns: ``{"name", "type"}`` dicts.
+ """
+ counts = Counter(col["type"] for col in columns)
+ ordered = sorted(counts.items(), key=lambda item: (-item[1], item[0]))
+ if len(ordered) <= _SCHEMA_TYPE_HISTOGRAM_TOP_K:
+ return dict(ordered)
+ histogram = dict(ordered[:_SCHEMA_TYPE_HISTOGRAM_TOP_K])
+ histogram[_SCHEMA_TYPE_HISTOGRAM_OTHER_KEY] = sum(
+ count for _, count in ordered[_SCHEMA_TYPE_HISTOGRAM_TOP_K:]
+ )
+ return histogram
+
+
+def build_schema_result(
+ columns: Sequence[dict[str, str]],
+ *,
+ max_columns: int,
+ max_result_bytes: int,
+ name_contains: str | None = None,
+) -> str:
+ """
+ Render a table's columns as a bounded JSON tool result.
+
+ Column names are the information an agent needs to write SQL, so unlike
query rows they cannot
+ simply be dropped: above ``max_columns`` (or the byte budget) the full
list is replaced by a
+ summary -- count, a type histogram, and a sample of columns -- that names
``name_contains`` as
+ the way to retrieve specific columns. ``columns`` is assumed to hold
distinct names (a table's
+ introspected columns are unique by construction).
+
+ :param columns: ``{"name", "type"}`` dicts in table order.
+ :param max_columns: Column count above which a summary replaces the full
list.
+ :param max_result_bytes: Budget for the serialized result.
+ :param name_contains: Case-insensitive substring. When given (and
non-empty), only matching
+ columns are considered and the value is echoed back so a filtered
subset is never mistaken
+ for the whole table.
+ """
+ name_contains = name_contains or None
+ total_columns = len(columns)
+ if name_contains is not None:
+ needle = name_contains.casefold()
+ selected = [col for col in columns if needle in col["name"].casefold()]
+ else:
+ selected = list(columns)
+
+ if name_contains is not None and not selected:
+ plural = "" if total_columns == 1 else "s"
+ return _dumps(
+ {
+ "columns": [],
+ "column_count": 0,
+ "name_contains": name_contains,
+ "hint": (
+ f"No columns match name_contains={name_contains!r}. Call
get_schema without "
+ f"name_contains to list all {total_columns}
column{plural}."
+ ),
+ }
+ )
+
+ full: dict[str, Any] = {"columns": selected, "column_count": len(selected)}
+ if name_contains is not None:
+ full["name_contains"] = name_contains
+ full["total_columns"] = total_columns
+ if len(selected) <= max_columns and _size(full) <= max_result_bytes:
+ return _dumps(full)
+
+ return _summarize_schema(
+ selected,
+ total_columns=total_columns,
+ max_columns=max_columns,
+ max_result_bytes=max_result_bytes,
+ name_contains=name_contains,
+ )
+
+
+def _summarize_schema(
+ selected: list[dict[str, str]],
+ *,
+ total_columns: int,
+ max_columns: int,
+ max_result_bytes: int,
+ name_contains: str | None,
+) -> str:
+ """Build the bounded summary returned when the full column list does not
fit."""
+ # Count is checked before bytes: it is the cheaper, more explainable bound
and the one the
+ # issue is about. A bytes-only truncation then means "count fits but
names/types are
+ # pathologically long", a distinct and rarer signal.
+ truncated_by = "max_columns" if len(selected) > max_columns else
"max_result_bytes"
+ hint = (
+ f"This filter matches {len(selected)} columns, more than can be
returned at once. "
+ "Use a more specific name_contains substring to narrow to the columns
you need."
+ if name_contains is not None
+ else (
+ f"This table has {len(selected)} columns, more than can be
returned at once. Call "
+ "get_schema again with a name_contains substring to return only
the matching columns."
+ )
+ )
+
+ output: dict[str, Any] = {
+ "column_count": len(selected),
+ "truncated": True,
+ "truncated_by": truncated_by,
+ "hint": hint,
+ }
+ if name_contains is not None:
+ output["name_contains"] = name_contains
+ output["total_columns"] = total_columns
+
+ # The core above is the guaranteed-useful payload. Add the histogram and a
sample of columns
+ # only while they fit, reserving the core first -- the same
contiguous-prefix accounting
+ # build_query_result uses for rows -- so a pathologically small budget
still returns the core.
+ output["type_histogram"] = _build_type_histogram(selected)
+ output["sample_columns"] = []
+ if _size(output) > max_result_bytes:
+ del output["type_histogram"]
+ return _dumps(output)
Review Comment:
The branch now drops the histogram and falls through to the sample loop
instead of returning early, so a wide struct type no longer strips the preview
off every other column. Added a test that forces the histogram over budget and
asserts the sample is still filled.
##########
providers/common/ai/src/airflow/providers/common/ai/utils/query_results.py:
##########
@@ -155,3 +188,136 @@ def build_query_result(
f"than paging through the result."
)
return _dumps(output)
+
+
+def _build_type_histogram(columns: Sequence[dict[str, str]]) -> dict[str, int]:
+ """
+ Count columns per type, capped to the most common
``_SCHEMA_TYPE_HISTOGRAM_TOP_K`` types.
+
+ Ordered by ``(-count, type)`` so the output is content-stable
(deterministic for byte
+ accounting and tests, not merely input-ordered). The long tail is folded
into a single
+ aggregate entry rather than listed, so parametrized types cannot inflate
the key count.
+
+ :param columns: ``{"name", "type"}`` dicts.
+ """
+ counts = Counter(col["type"] for col in columns)
+ ordered = sorted(counts.items(), key=lambda item: (-item[1], item[0]))
+ if len(ordered) <= _SCHEMA_TYPE_HISTOGRAM_TOP_K:
+ return dict(ordered)
+ histogram = dict(ordered[:_SCHEMA_TYPE_HISTOGRAM_TOP_K])
+ histogram[_SCHEMA_TYPE_HISTOGRAM_OTHER_KEY] = sum(
+ count for _, count in ordered[_SCHEMA_TYPE_HISTOGRAM_TOP_K:]
+ )
+ return histogram
+
+
+def build_schema_result(
+ columns: Sequence[dict[str, str]],
+ *,
+ max_columns: int,
+ max_result_bytes: int,
+ name_contains: str | None = None,
+) -> str:
+ """
+ Render a table's columns as a bounded JSON tool result.
+
+ Column names are the information an agent needs to write SQL, so unlike
query rows they cannot
+ simply be dropped: above ``max_columns`` (or the byte budget) the full
list is replaced by a
+ summary -- count, a type histogram, and a sample of columns -- that names
``name_contains`` as
+ the way to retrieve specific columns. ``columns`` is assumed to hold
distinct names (a table's
+ introspected columns are unique by construction).
+
+ :param columns: ``{"name", "type"}`` dicts in table order.
+ :param max_columns: Column count above which a summary replaces the full
list.
+ :param max_result_bytes: Budget for the serialized result.
+ :param name_contains: Case-insensitive substring. When given (and
non-empty), only matching
+ columns are considered and the value is echoed back so a filtered
subset is never mistaken
+ for the whole table.
+ """
+ name_contains = name_contains or None
+ total_columns = len(columns)
+ if name_contains is not None:
+ needle = name_contains.casefold()
+ selected = [col for col in columns if needle in col["name"].casefold()]
+ else:
+ selected = list(columns)
+
+ if name_contains is not None and not selected:
+ plural = "" if total_columns == 1 else "s"
+ return _dumps(
+ {
+ "columns": [],
+ "column_count": 0,
+ "name_contains": name_contains,
+ "hint": (
+ f"No columns match name_contains={name_contains!r}. Call
get_schema without "
+ f"name_contains to list all {total_columns}
column{plural}."
Review Comment:
The no-match result now carries `total_columns`, and when `total_columns` is
over `max_columns` it suggests a different substring instead of "call without
name_contains" (which would bounce back to this same summary). Added a test for
the wide case.
##########
providers/common/ai/src/airflow/providers/common/ai/utils/query_results.py:
##########
@@ -155,3 +188,136 @@ def build_query_result(
f"than paging through the result."
)
return _dumps(output)
+
+
+def _build_type_histogram(columns: Sequence[dict[str, str]]) -> dict[str, int]:
+ """
+ Count columns per type, capped to the most common
``_SCHEMA_TYPE_HISTOGRAM_TOP_K`` types.
+
+ Ordered by ``(-count, type)`` so the output is content-stable
(deterministic for byte
+ accounting and tests, not merely input-ordered). The long tail is folded
into a single
+ aggregate entry rather than listed, so parametrized types cannot inflate
the key count.
+
+ :param columns: ``{"name", "type"}`` dicts.
+ """
+ counts = Counter(col["type"] for col in columns)
+ ordered = sorted(counts.items(), key=lambda item: (-item[1], item[0]))
+ if len(ordered) <= _SCHEMA_TYPE_HISTOGRAM_TOP_K:
+ return dict(ordered)
+ histogram = dict(ordered[:_SCHEMA_TYPE_HISTOGRAM_TOP_K])
+ histogram[_SCHEMA_TYPE_HISTOGRAM_OTHER_KEY] = sum(
+ count for _, count in ordered[_SCHEMA_TYPE_HISTOGRAM_TOP_K:]
+ )
+ return histogram
+
+
+def build_schema_result(
+ columns: Sequence[dict[str, str]],
+ *,
+ max_columns: int,
+ max_result_bytes: int,
+ name_contains: str | None = None,
+) -> str:
+ """
+ Render a table's columns as a bounded JSON tool result.
+
+ Column names are the information an agent needs to write SQL, so unlike
query rows they cannot
+ simply be dropped: above ``max_columns`` (or the byte budget) the full
list is replaced by a
+ summary -- count, a type histogram, and a sample of columns -- that names
``name_contains`` as
+ the way to retrieve specific columns. ``columns`` is assumed to hold
distinct names (a table's
+ introspected columns are unique by construction).
+
+ :param columns: ``{"name", "type"}`` dicts in table order.
+ :param max_columns: Column count above which a summary replaces the full
list.
+ :param max_result_bytes: Budget for the serialized result.
+ :param name_contains: Case-insensitive substring. When given (and
non-empty), only matching
+ columns are considered and the value is echoed back so a filtered
subset is never mistaken
+ for the whole table.
+ """
+ name_contains = name_contains or None
+ total_columns = len(columns)
+ if name_contains is not None:
+ needle = name_contains.casefold()
+ selected = [col for col in columns if needle in col["name"].casefold()]
+ else:
+ selected = list(columns)
+
+ if name_contains is not None and not selected:
+ plural = "" if total_columns == 1 else "s"
+ return _dumps(
+ {
+ "columns": [],
+ "column_count": 0,
+ "name_contains": name_contains,
+ "hint": (
+ f"No columns match name_contains={name_contains!r}. Call
get_schema without "
+ f"name_contains to list all {total_columns}
column{plural}."
+ ),
+ }
+ )
+
+ full: dict[str, Any] = {"columns": selected, "column_count": len(selected)}
+ if name_contains is not None:
+ full["name_contains"] = name_contains
+ full["total_columns"] = total_columns
+ if len(selected) <= max_columns and _size(full) <= max_result_bytes:
+ return _dumps(full)
+
+ return _summarize_schema(
+ selected,
+ total_columns=total_columns,
+ max_columns=max_columns,
+ max_result_bytes=max_result_bytes,
+ name_contains=name_contains,
+ )
+
+
+def _summarize_schema(
+ selected: list[dict[str, str]],
+ *,
+ total_columns: int,
+ max_columns: int,
+ max_result_bytes: int,
+ name_contains: str | None,
+) -> str:
+ """Build the bounded summary returned when the full column list does not
fit."""
+ # Count is checked before bytes: it is the cheaper, more explainable bound
and the one the
+ # issue is about. A bytes-only truncation then means "count fits but
names/types are
+ # pathologically long", a distinct and rarer signal.
+ truncated_by = "max_columns" if len(selected) > max_columns else
"max_result_bytes"
+ hint = (
+ f"This filter matches {len(selected)} columns, more than can be
returned at once. "
+ "Use a more specific name_contains substring to narrow to the columns
you need."
+ if name_contains is not None
+ else (
+ f"This table has {len(selected)} columns, more than can be
returned at once. Call "
Review Comment:
The hint is now keyed on `truncated_by` (over the `max_columns` limit vs did
not fit `max_result_bytes`) and pluralizes, so a single 70 KB column name no
longer reads "This table has 1 columns". Added
`test_summary_hint_is_worded_from_the_limit_it_hit` to pin both branches.
##########
providers/common/ai/src/airflow/providers/common/ai/utils/query_results.py:
##########
@@ -155,3 +188,136 @@ def build_query_result(
f"than paging through the result."
)
return _dumps(output)
+
+
+def _build_type_histogram(columns: Sequence[dict[str, str]]) -> dict[str, int]:
+ """
+ Count columns per type, capped to the most common
``_SCHEMA_TYPE_HISTOGRAM_TOP_K`` types.
+
+ Ordered by ``(-count, type)`` so the output is content-stable
(deterministic for byte
+ accounting and tests, not merely input-ordered). The long tail is folded
into a single
+ aggregate entry rather than listed, so parametrized types cannot inflate
the key count.
+
+ :param columns: ``{"name", "type"}`` dicts.
+ """
+ counts = Counter(col["type"] for col in columns)
+ ordered = sorted(counts.items(), key=lambda item: (-item[1], item[0]))
+ if len(ordered) <= _SCHEMA_TYPE_HISTOGRAM_TOP_K:
+ return dict(ordered)
+ histogram = dict(ordered[:_SCHEMA_TYPE_HISTOGRAM_TOP_K])
+ histogram[_SCHEMA_TYPE_HISTOGRAM_OTHER_KEY] = sum(
+ count for _, count in ordered[_SCHEMA_TYPE_HISTOGRAM_TOP_K:]
+ )
+ return histogram
+
+
+def build_schema_result(
+ columns: Sequence[dict[str, str]],
+ *,
+ max_columns: int,
+ max_result_bytes: int,
+ name_contains: str | None = None,
+) -> str:
+ """
+ Render a table's columns as a bounded JSON tool result.
+
+ Column names are the information an agent needs to write SQL, so unlike
query rows they cannot
+ simply be dropped: above ``max_columns`` (or the byte budget) the full
list is replaced by a
+ summary -- count, a type histogram, and a sample of columns -- that names
``name_contains`` as
+ the way to retrieve specific columns. ``columns`` is assumed to hold
distinct names (a table's
+ introspected columns are unique by construction).
+
+ :param columns: ``{"name", "type"}`` dicts in table order.
+ :param max_columns: Column count above which a summary replaces the full
list.
+ :param max_result_bytes: Budget for the serialized result.
+ :param name_contains: Case-insensitive substring. When given (and
non-empty), only matching
+ columns are considered and the value is echoed back so a filtered
subset is never mistaken
+ for the whole table.
+ """
+ name_contains = name_contains or None
+ total_columns = len(columns)
+ if name_contains is not None:
+ needle = name_contains.casefold()
+ selected = [col for col in columns if needle in col["name"].casefold()]
+ else:
+ selected = list(columns)
+
+ if name_contains is not None and not selected:
+ plural = "" if total_columns == 1 else "s"
+ return _dumps(
+ {
+ "columns": [],
+ "column_count": 0,
+ "name_contains": name_contains,
+ "hint": (
+ f"No columns match name_contains={name_contains!r}. Call
get_schema without "
+ f"name_contains to list all {total_columns}
column{plural}."
+ ),
+ }
+ )
+
+ full: dict[str, Any] = {"columns": selected, "column_count": len(selected)}
+ if name_contains is not None:
+ full["name_contains"] = name_contains
+ full["total_columns"] = total_columns
+ if len(selected) <= max_columns and _size(full) <= max_result_bytes:
+ return _dumps(full)
+
+ return _summarize_schema(
+ selected,
+ total_columns=total_columns,
+ max_columns=max_columns,
+ max_result_bytes=max_result_bytes,
+ name_contains=name_contains,
+ )
+
+
+def _summarize_schema(
+ selected: list[dict[str, str]],
+ *,
+ total_columns: int,
+ max_columns: int,
+ max_result_bytes: int,
+ name_contains: str | None,
+) -> str:
+ """Build the bounded summary returned when the full column list does not
fit."""
+ # Count is checked before bytes: it is the cheaper, more explainable bound
and the one the
+ # issue is about. A bytes-only truncation then means "count fits but
names/types are
Review Comment:
Removed that phrasing from the comment.
##########
providers/common/ai/docs/changelog.rst:
##########
@@ -39,6 +39,16 @@ Changelog
attempt is checked and counted on its own, unchanged. See :ref:`the
cross-attempt usage
budget <agent-usage-budget>`.
+.. note::
+ ``get_schema`` on ``SQLToolset`` and ``DataFusionToolset`` now returns a
JSON object
+ ``{"columns": [{"name", "type"}, ...], "column_count": N}`` instead of a
bare JSON array of
+ columns. The tool also accepts an optional ``name_contains`` substring
filter, and on a table
+ with more columns than ``max_columns`` (default 100), or one whose
serialized columns exceed
+ ``max_result_bytes``, it returns a bounded summary (``column_count``, a
``type_histogram`` and a
+ ``sample_columns`` preview, with ``truncated``, ``truncated_by`` and a
``hint``) in place of the
+ full list. Update any system prompt or direct ``call_tool("get_schema",
...)`` caller that read
+ the old top-level array to read ``result["columns"]`` instead.
Review Comment:
The note now says to check `truncated` first, then read `columns` on a full
result or `sample_columns` on a summary, and states that a summary has no
`columns` key so `result["columns"]` raises `KeyError` on a table wide enough
to be summarized.
##########
providers/common/ai/docs/toolsets/sql.rst:
##########
@@ -205,6 +206,8 @@ Parameters
transferred is its own call. See :ref:`bounded-query-results`.
- ``max_result_bytes``: Budget for the serialized ``query`` result. Default 64
KiB.
Review Comment:
Updated the wording in `sql.rst`, `datafusion.rst`, and both
`max_result_bytes` docstrings to note it also triggers the `get_schema`
summary, with a cross-reference to the bounded-schema-results section.
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
