You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by hu...@apache.org on 2026/10/20 22:54:25 UTC
(superset) branch hughhhh/pfm-6-explore-indicator updated: fix(samples): skip the unbounded COUNT(*) on partitioned datasets
This is an automated email from the ASF dual-hosted git repository.
hughhhh pushed a commit to branch hughhhh/pfm-6-explore-indicator
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/hughhhh/pfm-6-explore-indicator by this push:
new 1764827c2b3 fix(samples): skip the unbounded COUNT(*) on partitioned datasets
1764827c2b3 is described below
commit 1764827c2b35f8a0779f17ebf31d280d6e3e8926
Author: Hugh A Miles II <hu...@gmail.com>
AuthorDate: Tue Oct 20 23:39:00 2026 +0100
fix(samples): skip the unbounded COUNT(*) on partitioned datasets
Samples (Explore data pane and drill-to-detail, both via
/datasource/samples) ran a separate COUNT(*) with no row limit and,
unless the chart happened to filter the mapped column, no partition
predicate. On a 40B-row Impala/Hive table that is a full scan that runs
until it times out.
For datasets with a partition column (and PARTITION_FILTER_MAPPING on),
skip the count. The rows query, which is already bounded, asks for one
probe row past the page instead: if it comes back, the probe is trimmed
and the total is reported as a lower bound (total_count_is_lower_bound);
if not, the page is the last one and the total is exact. Other datasets
keep the exact count.
The frontend shows a lower-bound total as "N+ rows", and drill-to-detail
counts one row past it so the next page stays reachable.
Co-Authored-By: Claude Opus 5.5 <no...@anthropic.com>
---
.../Chart/DrillDetail/DrillDetailPane.test.tsx | 48 ++++
.../Chart/DrillDetail/DrillDetailPane.tsx | 13 +-
.../Chart/DrillDetail/DrillDetailTableControls.tsx | 8 +-
.../src/components/Chart/DrillDetail/types.ts | 3 +
.../RowCountLabel/RowCountLabel.test.tsx | 5 +
.../src/components/RowCountLabel/index.tsx | 8 +-
.../components/DataTableControls.tsx | 2 +
.../DataTablesPane/components/SamplesPane.tsx | 11 +
.../DataTablesPane/test/SamplesPane.test.tsx | 30 +++
.../src/explore/components/DataTablesPane/types.ts | 2 +
superset/views/datasource/utils.py | 146 +++++++++---
tests/unit_tests/views/datasource/test_utils.py | 252 +++++++++++++++++++++
12 files changed, 496 insertions(+), 32 deletions(-)
diff --git a/superset-frontend/src/components/Chart/DrillDetail/DrillDetailPane.test.tsx b/superset-frontend/src/components/Chart/DrillDetail/DrillDetailPane.test.tsx
index e066c133c73..5d8e411eb89 100644
--- a/superset-frontend/src/components/Chart/DrillDetail/DrillDetailPane.test.tsx
+++ b/superset-frontend/src/components/Chart/DrillDetail/DrillDetailPane.test.tsx
@@ -171,6 +171,34 @@ const fetchWithTwoPages = () => {
});
};
+// A partitioned dataset: no COUNT(*), only "a full page came back and more
+// rows exist past it".
+const fetchWithLowerBoundTotal = () => {
+ setupDatasetEndpoint();
+ fetchMock.post(SAMPLES_ENDPOINT, {
+ result: {
+ total_count: 50,
+ total_count_is_lower_bound: true,
+ data: Array.from({ length: 50 }, (_, i) => ({
+ year: 1950 + i,
+ na_sales: i,
+ eu_sales: i,
+ })),
+ colnames: ['year', 'na_sales', 'eu_sales'],
+ coltypes: [0, 0, 0],
+ },
+ });
+ fetchMock.post(SAMPLES_ENDPOINT_PAGE_2, {
+ result: {
+ total_count: 51,
+ total_count_is_lower_bound: false,
+ data: [{ year: 2020, na_sales: 42.42, eu_sales: 24.24 }],
+ colnames: ['year', 'na_sales', 'eu_sales'],
+ coltypes: [0, 0, 0],
+ },
+ });
+};
+
afterEach(() => {
fetchMock.clearHistory().removeRoutes();
supersetGetCache.clear();
@@ -488,3 +516,23 @@ test('resets to page 1 and refetches on reload', async () => {
expect(screen.queryByText('42.42')).not.toBeInTheDocument();
expect(fetchMock.callHistory.calls(SAMPLES_ENDPOINT)).toHaveLength(2);
});
+
+test('pages past a lower-bound total and settles on the exact one', async () => {
+ jest.restoreAllMocks();
+ fetchWithLowerBoundTotal();
+ await waitForRender();
+
+ expect(await screen.findByText('50+ rows')).toBeInTheDocument();
+ // 50 rows on a 50-row page would be a single page; the lower bound has to
+ // keep page 2 reachable.
+ const pagination = await waitFor(() => {
+ const element = document.querySelector<HTMLElement>('.ant-pagination');
+ expect(element).toBeTruthy();
+ return element as HTMLElement;
+ });
+ await userEvent.click(within(pagination).getByTitle('2'));
+
+ expect(await screen.findByText('2020')).toBeInTheDocument();
+ expect(screen.getByText('51 rows')).toBeInTheDocument();
+ expect(fetchMock.callHistory.calls(SAMPLES_ENDPOINT_PAGE_2)).toHaveLength(1);
+});
diff --git a/superset-frontend/src/components/Chart/DrillDetail/DrillDetailPane.tsx b/superset-frontend/src/components/Chart/DrillDetail/DrillDetailPane.tsx
index ea41b55d941..d524a6fa69f 100644
--- a/superset-frontend/src/components/Chart/DrillDetail/DrillDetailPane.tsx
+++ b/superset-frontend/src/components/Chart/DrillDetail/DrillDetailPane.tsx
@@ -144,6 +144,13 @@ export default function DrillDetailPane({
return resultsPages.get(lastPageIndex.current);
}, [pageIndex, resultsPages]);
+ // A lower-bound total only covers the rows up to the current page, so count
+ // one row past it to keep the next page reachable.
+ const recordCount =
+ resultsPage && resultsPage.totalIsLowerBound
+ ? resultsPage.total + 1
+ : resultsPage?.total;
+
const mappedColumns: ColumnsType<DataType> = useMemo(
() =>
resultsPage?.colNames.map((column, index) => ({
@@ -327,6 +334,9 @@ export default function DrillDetailPane({
pageIndex,
{
total: response.total_count,
+ totalIsLowerBound: Boolean(
+ response.total_count_is_lower_bound,
+ ),
data: response.data,
colNames: ensureIsArray(response.colnames),
colTypes: ensureIsArray(response.coltypes),
@@ -397,7 +407,7 @@ export default function DrillDetailPane({
columns={mappedColumns}
size={TableSize.Small}
defaultPageSize={DEFAULT_PAGE_SIZE}
- recordCount={resultsPage?.total}
+ recordCount={recordCount}
usePagination
loading={isLoading}
onChange={pagination => {
@@ -426,6 +436,7 @@ export default function DrillDetailPane({
filters={filters}
setFilters={setFilters}
totalCount={resultsPage?.total}
+ totalCountIsLowerBound={resultsPage?.totalIsLowerBound}
loading={isLoading}
onReload={handleReload}
canDownload={canDownload}
diff --git a/superset-frontend/src/components/Chart/DrillDetail/DrillDetailTableControls.tsx b/superset-frontend/src/components/Chart/DrillDetail/DrillDetailTableControls.tsx
index 9771c3faf55..bf358c0d90d 100644
--- a/superset-frontend/src/components/Chart/DrillDetail/DrillDetailTableControls.tsx
+++ b/superset-frontend/src/components/Chart/DrillDetail/DrillDetailTableControls.tsx
@@ -37,6 +37,7 @@ export type TableControlsProps = {
filters: BinaryQueryObjectFilterClause[];
setFilters: (filters: BinaryQueryObjectFilterClause[]) => void;
totalCount?: number;
+ totalCountIsLowerBound?: boolean;
loading: boolean;
onReload: () => void;
canDownload: boolean;
@@ -50,6 +51,7 @@ export default function TableControls({
filters,
setFilters,
totalCount,
+ totalCountIsLowerBound,
loading,
onReload,
canDownload,
@@ -137,7 +139,11 @@ export default function TableControls({
gap: ${theme.sizeUnit * 3}px;
`}
>
- <RowCountLabel loading={loading && !totalCount} rowcount={totalCount} />
+ <RowCountLabel
+ loading={loading && !totalCount}
+ rowcount={totalCount}
+ isLowerBound={totalCountIsLowerBound}
+ />
{canDownload && (
<DownloadDropdown
onDownloadCSV={onDownloadCSV}
diff --git a/superset-frontend/src/components/Chart/DrillDetail/types.ts b/superset-frontend/src/components/Chart/DrillDetail/types.ts
index f300e165a1a..5620bf88224 100644
--- a/superset-frontend/src/components/Chart/DrillDetail/types.ts
+++ b/superset-frontend/src/components/Chart/DrillDetail/types.ts
@@ -20,6 +20,9 @@ import { GenericDataType } from '@apache-superset/core/common';
export type ResultsPage = {
total: number;
+ // `total` counts the rows up to this page and more exist past it; the
+ // backend skipped an exact COUNT(*) that would have been a full-table scan.
+ totalIsLowerBound?: boolean;
data: Record<string, any>[];
colNames: string[];
colTypes: GenericDataType[];
diff --git a/superset-frontend/src/components/RowCountLabel/RowCountLabel.test.tsx b/superset-frontend/src/components/RowCountLabel/RowCountLabel.test.tsx
index cbc6278ddaf..da5a2b4bc30 100644
--- a/superset-frontend/src/components/RowCountLabel/RowCountLabel.test.tsx
+++ b/superset-frontend/src/components/RowCountLabel/RowCountLabel.test.tsx
@@ -69,6 +69,11 @@ test('RowCountLabel uses a caller-provided limitReachedMessage instead of the de
expect(tooltip).not.toHaveTextContent('The row limit set for the chart');
});
+test('RowCountLabel marks a lower-bound count with a plus', () => {
+ render(<RowCountLabel rowcount={1000} isLowerBound />);
+ expect(screen.getByText('1k+ rows')).toBeInTheDocument();
+});
+
test('RowCountLabel renders loading', async () => {
render(<RowCountLabel loading />);
const expectedText = 'Loading...';
diff --git a/superset-frontend/src/components/RowCountLabel/index.tsx b/superset-frontend/src/components/RowCountLabel/index.tsx
index 3523fc0cae4..e3c135306af 100644
--- a/superset-frontend/src/components/RowCountLabel/index.tsx
+++ b/superset-frontend/src/components/RowCountLabel/index.tsx
@@ -29,6 +29,9 @@ type RowCountLabelProps = {
// Overrides the default "chart" wording for panes (e.g. samples) where the
// limit reached isn't the chart's own row_limit.
limitReachedMessage?: React.ReactNode;
+ // The total is a floor rather than an exact count, e.g. samples from a
+ // partitioned table where counting every row would be a full scan.
+ isLowerBound?: boolean;
};
const limitReachedMsg = t(
@@ -42,11 +45,14 @@ export default function RowCountLabel(props: RowCountLabelProps) {
loading,
label,
limitReachedMessage,
+ isLowerBound = false,
} = props;
const limitReached = limit && rowcount >= limit;
const type =
limitReached || (rowcount === 0 && !loading) ? 'error' : 'default';
- const formattedRowCount = getNumberFormatter()(rowcount);
+ const formattedRowCount = `${getNumberFormatter()(rowcount)}${
+ isLowerBound ? '+' : ''
+ }`;
const labelText = (
<Label type={type} monospace>
{loading ? (
diff --git a/superset-frontend/src/explore/components/DataTablesPane/components/DataTableControls.tsx b/superset-frontend/src/explore/components/DataTablesPane/components/DataTableControls.tsx
index ac21f5ab68a..e9741701866 100644
--- a/superset-frontend/src/explore/components/DataTablesPane/components/DataTableControls.tsx
+++ b/superset-frontend/src/explore/components/DataTablesPane/components/DataTableControls.tsx
@@ -64,6 +64,7 @@ export const TableControls = ({
columnNames,
columnTypes,
rowcount,
+ rowcountIsLowerBound,
isLoading,
canDownload,
rowLimit,
@@ -125,6 +126,7 @@ export const TableControls = ({
limit={effectiveRowLimit ?? rowLimit}
limitReachedMessage={limitReachedMessage}
loading={isLoading}
+ isLowerBound={rowcountIsLowerBound}
/>
{canDownload && onDownloadCSV && onDownloadXLSX && (
<DownloadDropdown
diff --git a/superset-frontend/src/explore/components/DataTablesPane/components/SamplesPane.tsx b/superset-frontend/src/explore/components/DataTablesPane/components/SamplesPane.tsx
index 7437d5e742a..796b49381ed 100644
--- a/superset-frontend/src/explore/components/DataTablesPane/components/SamplesPane.tsx
+++ b/superset-frontend/src/explore/components/DataTablesPane/components/SamplesPane.tsx
@@ -71,6 +71,8 @@ export const SamplesPane = ({
const [coltypes, setColtypes] = useState<GenericDataType[]>([]);
const [isLoading, setIsLoading] = useState<boolean>(false);
const [rowcount, setRowCount] = useState<number>(0);
+ const [rowcountIsLowerBound, setRowcountIsLowerBound] =
+ useState<boolean>(false);
const [responseError, setResponseError] = useState<string>('');
const { gridHeight, measuredRef } = useGridHeight();
const datasourceId = useMemo(
@@ -117,6 +119,12 @@ export const SamplesPane = ({
// controls only render when there are rows, and a hardcoded 0 would
// label a populated table as "0 rows".
setRowCount(response?.rowcount ?? rows.length);
+ // Partitioned datasets skip the exact COUNT(*), which would scan the
+ // whole table; the backend flags instead that rows exist past this
+ // page. `total_count` itself is not needed here.
+ setRowcountIsLowerBound(
+ Boolean(response?.total_count_is_lower_bound),
+ );
setResponseError('');
cache.set(queryFormData, true);
if (queryForce) {
@@ -127,6 +135,7 @@ export const SamplesPane = ({
setData([]);
setColnames([]);
setColtypes([]);
+ setRowcountIsLowerBound(false);
setResponseError(`${error.name}: ${error.message}`);
})
.finally(() => {
@@ -161,6 +170,7 @@ export const SamplesPane = ({
columnNames={colnames}
columnTypes={coltypes}
rowcount={rowcount}
+ rowcountIsLowerBound={rowcountIsLowerBound}
datasourceId={datasourceId}
onInputChange={handleInputChange}
filterText={filterText}
@@ -202,6 +212,7 @@ export const SamplesPane = ({
columnNames={colnames}
columnTypes={coltypes}
rowcount={rowcount}
+ rowcountIsLowerBound={rowcountIsLowerBound}
datasourceId={datasourceId}
onInputChange={handleInputChange}
filterText={filterText}
diff --git a/superset-frontend/src/explore/components/DataTablesPane/test/SamplesPane.test.tsx b/superset-frontend/src/explore/components/DataTablesPane/test/SamplesPane.test.tsx
index 19839688e5d..feff7508976 100644
--- a/superset-frontend/src/explore/components/DataTablesPane/test/SamplesPane.test.tsx
+++ b/superset-frontend/src/explore/components/DataTablesPane/test/SamplesPane.test.tsx
@@ -81,6 +81,25 @@ describe('SamplesPane', () => {
},
);
+ // A partitioned dataset: the backend skipped COUNT(*) and reports that rows
+ // exist past this page instead of an exact total.
+ fetchMock.post(
+ 'end:/datasource/samples?force=false&datasource_type=table&datasource_id=39&per_page=100&page=1',
+ {
+ result: {
+ data: [
+ { __timestamp: 1230768000000, genre: 'Action' },
+ { __timestamp: 1230768000010, genre: 'Horror' },
+ ],
+ colnames: ['__timestamp', 'genre'],
+ coltypes: [2, 1],
+ rowcount: 2,
+ total_count: 2,
+ total_count_is_lower_bound: true,
+ },
+ },
+ );
+
const setForceQuery = jest.fn();
afterAll(() => {
@@ -161,4 +180,15 @@ describe('SamplesPane', () => {
expect(queryByText('0 rows')).not.toBeInTheDocument();
expect(queryByText('2 rows')).toBeVisible();
});
+
+ test('labels the row count as a lower bound when the total is not exact', async () => {
+ const props = createSamplesPaneProps({ datasourceId: 39 });
+ const { findByText, queryByText } = render(<SamplesPane {...props} />, {
+ useRedux: true,
+ });
+
+ expect(await findByText('Action')).toBeVisible();
+ expect(queryByText('2+ rows')).toBeVisible();
+ expect(queryByText('2 rows')).not.toBeInTheDocument();
+ });
});
diff --git a/superset-frontend/src/explore/components/DataTablesPane/types.ts b/superset-frontend/src/explore/components/DataTablesPane/types.ts
index 7b3be4399f5..8102ee2a98c 100644
--- a/superset-frontend/src/explore/components/DataTablesPane/types.ts
+++ b/superset-frontend/src/explore/components/DataTablesPane/types.ts
@@ -81,6 +81,8 @@ export interface TableControlsProps extends DrillControlsProps {
columnTypes: GenericDataType[];
isLoading: boolean;
rowcount: number;
+ // More rows exist past `rowcount`, which is then shown as "N+ rows".
+ rowcountIsLowerBound?: boolean;
canDownload: boolean;
rowLimit?: number;
rowLimitOptions?: { value: number; label: string }[];
diff --git a/superset/views/datasource/utils.py b/superset/views/datasource/utils.py
index 4f86e0d6d1c..3db443b92ee 100644
--- a/superset/views/datasource/utils.py
+++ b/superset/views/datasource/utils.py
@@ -22,6 +22,7 @@ from typing import Any, Iterable, Optional, TYPE_CHECKING
import pandas as pd
from flask import current_app as app
+from superset import is_feature_enabled
from superset.commands.dataset.exceptions import DatasetSamplesFailedError
from superset.common.chart_data import ChartDataResultType
from superset.common.query_context_factory import QueryContextFactory
@@ -101,6 +102,27 @@ def replace_verbose_with_column(
f["col"] = match
+def skips_count_star(datasource: Datasource) -> bool:
+ """
+ Whether Samples should skip its ``COUNT(*)`` for this datasource.
+
+ A dataset with a declared partition column is a partitioned warehouse table
+ (Impala/Hive in practice), often tens of billions of rows. The samples rows
+ query is bounded by its row limit, but the count is not: unless the request
+ happens to filter the mapped column, it is an unfiltered full-table scan,
+ which on those tables runs until it times out. For these datasets the total
+ is reported as a lower bound instead, derived from one extra probe row on
+ the bounded rows query, so Samples never issues a query that is both
+ unfiltered by partition and unbounded.
+
+ Gated on the feature flag like the rest of partition filter mapping, so a
+ deployment that never turned it on keeps the exact count.
+ """
+ return is_feature_enabled("PARTITION_FILTER_MAPPING") and bool(
+ getattr(datasource, "partition_column", None)
+ )
+
+
def get_samples( # pylint: disable=too-many-arguments
datasource_type: str,
datasource_id: int,
@@ -119,6 +141,18 @@ def get_samples( # pylint: disable=too-many-arguments
form_data = {"dashboardId": dashboard_id} if dashboard_id else None
limit_clause = get_limit_clause(page, per_page)
+ skip_count = skips_count_star(datasource)
+ # Without a count, one row past the page is what tells "there is another
+ # page" apart from "this was the last one".
+ query_limit_clause = (
+ {
+ **limit_clause,
+ "row_limit": (limit_clause["row_limit"] or app.config["SAMPLES_ROW_LIMIT"])
+ + 1,
+ }
+ if skip_count
+ else limit_clause
+ )
# todo(yongjie): Constructing count(*) and samples in the same query_context,
if payload is None:
@@ -128,7 +162,7 @@ def get_samples( # pylint: disable=too-many-arguments
"type": datasource.type,
"id": datasource.id,
},
- queries=[limit_clause],
+ queries=[query_limit_clause],
form_data=form_data,
result_type=ChartDataResultType.SAMPLES,
force=force,
@@ -145,31 +179,16 @@ def get_samples( # pylint: disable=too-many-arguments
"type": datasource.type,
"id": datasource.id,
},
- queries=[{**payload, **limit_clause}],
+ queries=[{**payload, **query_limit_clause}],
form_data=form_data,
result_type=ChartDataResultType.DRILL_DETAIL,
force=force,
)
- # constructing count(*) query
- count_star_metric = {
- "metrics": [
- {
- "expressionType": "SQL",
- "sqlExpression": "COUNT(*)",
- "label": "COUNT(*)",
- }
- ]
- }
- count_star_instance = QueryContextFactory().create(
- datasource={
- "type": datasource.type,
- "id": datasource.id,
- },
- queries=[{**payload, **count_star_metric} if payload else count_star_metric],
- form_data=form_data,
- result_type=ChartDataResultType.FULL,
- force=force,
+ count_star_instance = (
+ None
+ if skip_count
+ else _count_star_query_context(datasource, payload, form_data, force)
)
try:
@@ -177,18 +196,14 @@ def get_samples( # pylint: disable=too-many-arguments
# This prevents users with "can samples on Datasource" permission from
# reading samples from datasets they don't have access to.
samples_instance.raise_for_access()
- count_star_instance.raise_for_access()
-
- count_star_data = count_star_instance.get_payload()["queries"][0]
-
- if count_star_data.get("status") == QueryStatus.FAILED:
- raise DatasetSamplesFailedError(count_star_data.get("error"))
+ count_star_data = _get_count_star_data(count_star_instance)
engine_spec = datasource.database.db_engine_spec
row_offset = limit_clause["row_offset"]
row_limit = limit_clause["row_limit"]
- if not engine_spec.supports_offset and row_offset > 0:
+ via_cursor = not engine_spec.supports_offset and row_offset > 0
+ if via_cursor:
try:
sample_data = _fetch_samples_via_cursor(
datasource=datasource,
@@ -217,12 +232,85 @@ def get_samples( # pylint: disable=too-many-arguments
sample_data["page"] = page
sample_data["per_page"] = per_page
- sample_data["total_count"] = count_star_data["data"][0]["COUNT(*)"]
+ if skip_count:
+ _set_lower_bound_total(
+ sample_data, row_offset, row_limit, probed=not via_cursor
+ )
+ else:
+ sample_data["total_count"] = count_star_data["data"][0]["COUNT(*)"]
+ sample_data["total_count_is_lower_bound"] = False
return sample_data
except (IndexError, KeyError) as exc:
raise DatasetSamplesFailedError from exc
+def _count_star_query_context(
+ datasource: Datasource,
+ payload: SamplesPayloadSchema | None,
+ form_data: dict[str, Any] | None,
+ force: bool,
+) -> QueryContext:
+ count_star_metric = {
+ "metrics": [
+ {
+ "expressionType": "SQL",
+ "sqlExpression": "COUNT(*)",
+ "label": "COUNT(*)",
+ }
+ ]
+ }
+ return QueryContextFactory().create(
+ datasource={
+ "type": datasource.type,
+ "id": datasource.id,
+ },
+ queries=[{**payload, **count_star_metric} if payload else count_star_metric],
+ form_data=form_data,
+ result_type=ChartDataResultType.FULL,
+ force=force,
+ )
+
+
+def _get_count_star_data(
+ count_star_instance: QueryContext | None,
+) -> dict[str, Any]:
+ """Run the count query, or return an empty payload when it was skipped."""
+ if count_star_instance is None:
+ return {}
+
+ count_star_instance.raise_for_access()
+ count_star_data = count_star_instance.get_payload()["queries"][0]
+ if count_star_data.get("status") == QueryStatus.FAILED:
+ raise DatasetSamplesFailedError(count_star_data.get("error"))
+ return count_star_data
+
+
+def _set_lower_bound_total(
+ sample_data: dict[str, Any], row_offset: int, row_limit: int, probed: bool
+) -> None:
+ """
+ Report the total as "the rows up to and including this page".
+
+ The rows query asked for one probe row past the page. If it came back there
+ is at least one more page, so the probe is trimmed and the total flagged as
+ a lower bound; if it did not, this page is the last one and the total is
+ exact. The cursor path pages with an exact page size and has no probe row
+ (``probed=False``), so a full page is taken to mean there may be more.
+ """
+ data = sample_data.get("data") or []
+ page_size = row_limit or app.config["SAMPLES_ROW_LIMIT"]
+ if probed:
+ has_more = len(data) > page_size
+ if has_more:
+ data = data[:page_size]
+ sample_data["data"] = data
+ sample_data["rowcount"] = len(data)
+ else:
+ has_more = len(data) >= page_size
+ sample_data["total_count"] = row_offset + len(data)
+ sample_data["total_count_is_lower_bound"] = has_more
+
+
def _fetch_samples_via_cursor(
datasource: Datasource,
samples_instance: QueryContext,
diff --git a/tests/unit_tests/views/datasource/test_utils.py b/tests/unit_tests/views/datasource/test_utils.py
index 5c69ae6ba64..f63c2fb5ac7 100644
--- a/tests/unit_tests/views/datasource/test_utils.py
+++ b/tests/unit_tests/views/datasource/test_utils.py
@@ -394,3 +394,255 @@ def test_get_samples_cursor_path_unused_for_page_one(
)
datasource.database.db_engine_spec.fetch_data_with_cursor.assert_not_called()
+
+
+def _samples_ctx(rows: list[dict[str, int]]) -> MagicMock:
+ ctx = MagicMock()
+ ctx.get_payload.return_value = {
+ "queries": [
+ {
+ "data": rows,
+ "colnames": ["a"],
+ "coltypes": [],
+ "rowcount": len(rows),
+ "status": "success",
+ }
+ ]
+ }
+ return ctx
+
+
+def _count_ctx(count: int) -> MagicMock:
+ ctx = MagicMock()
+ ctx.get_payload.return_value = {
+ "queries": [{"data": [{"COUNT(*)": count}], "status": "success"}]
+ }
+ return ctx
+
+
+def _created_queries(qcf: MagicMock) -> list[dict[str, object]]:
+ return [
+ call.kwargs["queries"][0] for call in qcf.return_value.create.call_args_list
+ ]
+
+
+@pytest.mark.parametrize("payload", [None, {"filters": [], "columns": ["a"]}])
+def test_get_samples_partitioned_dataset_skips_count_star(
+ fake_datasource_factory: FakeDatasourceFactory,
+ payload: dict[str, object] | None,
+) -> None:
+ """
+ A dataset with a partition column never builds the COUNT(*) query: on a
+ 40B-row Impala/Hive table it is an unbounded full scan. The rows query asks
+ for one probe row past the page instead, and since the probe came back the
+ total is reported as a lower bound with the probe trimmed off.
+ """
+ from superset.views.datasource import utils
+
+ datasource = fake_datasource_factory(supports_offset=True)
+ datasource.partition_column = "ds"
+
+ with (
+ patch.object(utils, "is_feature_enabled", return_value=True),
+ patch.object(utils, "QueryContextFactory") as qcf,
+ ):
+ qcf.return_value.create.side_effect = [
+ _samples_ctx([{"a": i} for i in range(51)])
+ ]
+
+ result = utils.get_samples(
+ datasource_type="table",
+ datasource_id=1,
+ page=3,
+ per_page=50,
+ payload=payload, # type: ignore[arg-type]
+ datasource=datasource,
+ )
+
+ queries = _created_queries(qcf)
+ assert len(queries) == 1
+ assert "metrics" not in queries[0]
+ assert queries[0]["row_limit"] == 51
+ assert queries[0]["row_offset"] == 100
+ assert result["data"] == [{"a": i} for i in range(50)]
+ assert result["rowcount"] == 50
+ assert result["total_count"] == 150
+ assert result["total_count_is_lower_bound"] is True
+
+
+def test_get_samples_partitioned_dataset_last_page_is_exact(
+ fake_datasource_factory: FakeDatasourceFactory,
+) -> None:
+ """
+ No probe row means this page is the last one, so offset plus the rows on it
+ is the exact total.
+ """
+ from superset.views.datasource import utils
+
+ datasource = fake_datasource_factory(supports_offset=True)
+ datasource.partition_column = "ds"
+
+ with (
+ patch.object(utils, "is_feature_enabled", return_value=True),
+ patch.object(utils, "QueryContextFactory") as qcf,
+ ):
+ qcf.return_value.create.side_effect = [_samples_ctx([{"a": 1}] * 7)]
+
+ result = utils.get_samples(
+ datasource_type="table",
+ datasource_id=1,
+ page=2,
+ per_page=50,
+ datasource=datasource,
+ )
+
+ assert len(result["data"]) == 7
+ assert result["total_count"] == 57
+ assert result["total_count_is_lower_bound"] is False
+
+
+def test_get_samples_partitioned_dataset_default_page_size_probes(
+ fake_datasource_factory: FakeDatasourceFactory,
+) -> None:
+ """
+ ``per_page=0`` resolves to the configured SAMPLES_ROW_LIMIT downstream; the
+ probe has to be one past that, not one past zero.
+ """
+ from flask import current_app
+
+ from superset.views.datasource import utils
+
+ datasource = fake_datasource_factory(supports_offset=True)
+ datasource.partition_column = "ds"
+ samples_row_limit = current_app.config["SAMPLES_ROW_LIMIT"]
+
+ with (
+ patch.object(utils, "is_feature_enabled", return_value=True),
+ patch.object(utils, "QueryContextFactory") as qcf,
+ ):
+ qcf.return_value.create.side_effect = [_samples_ctx([{"a": 1}] * 3)]
+
+ result = utils.get_samples(
+ datasource_type="table",
+ datasource_id=1,
+ page=1,
+ per_page=0,
+ datasource=datasource,
+ )
+
+ assert _created_queries(qcf)[0]["row_limit"] == samples_row_limit + 1
+ assert result["total_count"] == 3
+ assert result["total_count_is_lower_bound"] is False
+
+
+def test_get_samples_partitioned_dataset_cursor_path_full_page_is_lower_bound(
+ fake_datasource_factory: FakeDatasourceFactory,
+) -> None:
+ """
+ The cursor path pages by exact page size and carries no probe row, so a
+ full page is reported as a lower bound; there is still no COUNT(*).
+ """
+ from superset.views.datasource import utils
+
+ datasource = fake_datasource_factory(supports_offset=False)
+ datasource.partition_column = "ds"
+ datasource.database.db_engine_spec.fetch_data_with_cursor.return_value = (
+ [(i,) for i in range(50)],
+ ["a"],
+ )
+
+ with (
+ patch.object(utils, "is_feature_enabled", return_value=True),
+ patch.object(utils, "QueryContextFactory") as qcf,
+ patch.object(utils, "extract_dataframe_dtypes", return_value=[]),
+ ):
+ samples_ctx = MagicMock()
+ samples_ctx.datasource.get_query_str.return_value = "SELECT a FROM t"
+ qcf.return_value.create.side_effect = [samples_ctx]
+
+ result = utils.get_samples(
+ datasource_type="table",
+ datasource_id=1,
+ page=2,
+ per_page=50,
+ datasource=datasource,
+ )
+
+ assert len(_created_queries(qcf)) == 1
+ fetch = datasource.database.db_engine_spec.fetch_data_with_cursor
+ assert fetch.call_args.kwargs["page_index"] == 1
+ assert fetch.call_args.kwargs["page_size"] == 50
+ assert result["total_count"] == 100
+ assert result["total_count_is_lower_bound"] is True
+
+
+def test_get_samples_partitioned_dataset_samples_failure_still_raises(
+ fake_datasource_factory: FakeDatasourceFactory,
+) -> None:
+ """With no count query there is no count cache entry to evict."""
+ from superset.commands.dataset.exceptions import DatasetSamplesFailedError
+ from superset.views.datasource import utils
+
+ datasource = fake_datasource_factory(supports_offset=True)
+ datasource.partition_column = "ds"
+ samples_ctx = MagicMock()
+ samples_ctx.get_payload.return_value = {
+ "queries": [{"status": "failed", "error": "timed out"}]
+ }
+
+ with (
+ patch.object(utils, "is_feature_enabled", return_value=True),
+ patch.object(utils, "QueryContextFactory") as qcf,
+ ):
+ qcf.return_value.create.side_effect = [samples_ctx]
+
+ with pytest.raises(DatasetSamplesFailedError, match="timed out"):
+ utils.get_samples(
+ datasource_type="table",
+ datasource_id=1,
+ datasource=datasource,
+ )
+
+
+@pytest.mark.parametrize(
+ ("flag_enabled", "partition_column"),
+ [(False, "ds"), (True, None), (True, "")],
+)
+def test_get_samples_keeps_exact_count_without_partition_mapping(
+ fake_datasource_factory: FakeDatasourceFactory,
+ flag_enabled: bool,
+ partition_column: str | None,
+) -> None:
+ """
+ Datasets without a partition column, or deployments with the feature off,
+ keep the exact COUNT(*) and the un-probed row limit.
+ """
+ from superset.views.datasource import utils
+
+ datasource = fake_datasource_factory(supports_offset=True)
+ datasource.partition_column = partition_column
+
+ with (
+ patch.object(utils, "is_feature_enabled", return_value=flag_enabled),
+ patch.object(utils, "QueryContextFactory") as qcf,
+ ):
+ qcf.return_value.create.side_effect = [
+ _samples_ctx([{"a": 1}] * 50),
+ _count_ctx(1234),
+ ]
+
+ result = utils.get_samples(
+ datasource_type="table",
+ datasource_id=1,
+ page=1,
+ per_page=50,
+ datasource=datasource,
+ )
+
+ queries = _created_queries(qcf)
+ assert len(queries) == 2
+ assert queries[0]["row_limit"] == 50
+ assert queries[1]["metrics"][0]["sqlExpression"] == "COUNT(*)" # type: ignore[index]
+ assert len(result["data"]) == 50
+ assert result["total_count"] == 1234
+ assert result["total_count_is_lower_bound"] is False