kaxil commented on code in PR #74017:
URL: https://github.com/apache/airflow/pull/74017#discussion_r4158785074
##########
providers/common/ai/tests/unit/common/ai/utils/test_query_results.py:
##########
@@ -169,3 +173,141 @@ def
test_a_secret_that_json_escapes_is_masked_in_the_rows(register_secret):
result = _build(["user", "password"], [["admin", secret]])
assert result["rows"] == [["admin", "***"]]
+
+
+def _schema(columns, *, max_columns=100, max_result_bytes=65_536,
name_contains=None) -> dict:
+ return json.loads(
+ build_schema_result(
+ columns,
+ max_columns=max_columns,
+ max_result_bytes=max_result_bytes,
+ name_contains=name_contains,
+ )
+ )
+
+
+def _cols(n: int, *, type_: str = "VARCHAR", prefix: str = "col") ->
list[dict[str, str]]:
+ return [{"name": f"{prefix}_{i}", "type": type_} for i in range(n)]
+
+
+class TestSchemaResult:
+ def test_small_table_returns_every_column(self):
+ cols = _cols(3)
+ assert _schema(cols) == {"columns": cols, "column_count": 3}
+
+ def test_name_contains_filters_case_insensitively(self):
+ cols = [
+ {"name": "CustomerId", "type": "INT"},
+ {"name": "amount", "type": "NUMERIC"},
+ {"name": "customer_name", "type": "VARCHAR"},
+ ]
+ data = _schema(cols, name_contains="customer")
+ assert data["columns"] == [cols[0], cols[2]]
+ assert data["column_count"] == 2
+ assert data["name_contains"] == "customer"
+ assert data["total_columns"] == 3
+ assert "truncated" not in data
+
+ def test_empty_name_contains_is_treated_as_no_filter(self):
+ cols = _cols(3)
+ assert _schema(cols, name_contains="") == {"columns": cols,
"column_count": 3}
+
+ def test_name_contains_with_no_matches_guides_without_erroring(self):
+ data = _schema(_cols(5), name_contains="zzz")
+ assert data["columns"] == []
+ assert data["column_count"] == 0
+ assert data["name_contains"] == "zzz"
+ assert data["total_columns"] == 5
+ assert "error" not in data
+ assert "all 5 columns" in data["hint"]
+
+ def test_no_match_on_a_wide_table_does_not_point_back_at_the_filter(self):
+ """When the full list would itself summarize, the hint must not say
'call without it'."""
+ data = _schema(_cols(300), max_columns=100, name_contains="zzz")
+ assert data["columns"] == []
+ assert data["total_columns"] == 300
+ assert "without name_contains" not in data["hint"]
+ assert "different substring" in data["hint"]
+
+ def test_too_many_columns_are_summarized_not_listed(self):
Review Comment:
Nothing here builds a table with exactly `max_columns` columns, so changing
`<=` to `<` in `build_schema_result` (or `>` to `>=` for `truncated_by`) would
still pass. A 100-column case that asserts the full list comes back and a
101-column case that asserts `truncated_by == "max_columns"` would pin the
boundary the docs describe.
##########
providers/common/ai/src/airflow/providers/common/ai/utils/query_results.py:
##########
@@ -155,3 +193,148 @@ 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"
+ # Do not tell the agent to call without name_contains when the full
list would itself be
+ # summarized (total_columns > max_columns) -- that would bounce it
back to this filter.
+ if total_columns > max_columns:
+ advice = (
+ f"No columns match name_contains={name_contains!r}. Try a
different substring "
+ f"(the table has {total_columns} column{plural}, too many to
list in full)."
+ )
+ else:
+ advice = (
+ f"No columns match name_contains={name_contains!r}. Call
get_schema without "
+ f"name_contains to list all {total_columns} column{plural}."
+ )
+ return _dumps(
+ {
+ "columns": [],
+ "column_count": 0,
+ "name_contains": name_contains,
+ "total_columns": total_columns,
+ "hint": advice,
+ }
+ )
+
+ 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. A bytes-only
+ # truncation then means "count fits but names/types are pathologically
long", a rarer signal.
+ n = len(selected)
+ plural = "" if n == 1 else "s"
+ truncated_by = "max_columns" if n > max_columns else "max_result_bytes"
+ subject = f"name_contains={name_contains!r} matched" if name_contains is
not None else "This table has"
+ if truncated_by == "max_columns":
+ reason = f"{subject} {n} column{plural}, more than the max_columns
limit of {max_columns}."
+ else:
+ reason = f"{subject} {n} column{plural} that did not fit
max_result_bytes ({max_result_bytes})."
+ move = (
+ "Use a more specific name_contains substring to narrow to the columns
you need."
+ if name_contains is not None
+ else "Call get_schema with a name_contains substring to return only
the columns you need."
+ )
+
+ output: dict[str, Any] = {
+ "column_count": n,
+ "truncated": True,
+ "truncated_by": truncated_by,
+ "hint": f"{reason} {move}",
+ }
+ 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
small sample of
+ # columns only while they fit, reserving the core first. If the histogram
alone would not fit,
+ # drop it but still fill the sample from what remains, so a wide struct
type cannot strip the
+ # preview off every other column. 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"]
+ budget = max_result_bytes - _size(output)
+ for col in selected[:_SCHEMA_SAMPLE_SIZE]:
+ cost = _size(col) + (1 if output["sample_columns"] else 0)
+ if cost > budget:
+ break
Review Comment:
Following up on the struct-column case from the last round: dropping the
histogram fixed the early return, but this `break` still empties the preview
when the wide column is in the first 20. In a quick probe with one ~80 KB
struct type among 150 VARCHAR columns at the default budget, the struct at
position 0 gives `sample_columns: []` and at position 5 gives 5 entries.
`continue` (or keeping the name with an abbreviated type) would keep filling
past it. The new histogram test uses many medium-sized types, so it doesn't hit
this.
--
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]
