From ec97c6c3bdf6b92926d458f9cf45914422bc1116 Mon Sep 17 00:00:00 2001 From: Shaurya Chaturvedi Date: Tue, 22 Jul 2025 09:10:37 -0700 Subject: [PATCH 001/167] [timeseries] Introducing Apache ECharts timeseries results visualization in Controller UI (#16390) Co-authored-by: Shaurya Chaturvedi --- .../app/components/Query/MetricStatsTable.tsx | 212 +++++++++++++++ .../app/components/Query/TimeseriesChart.tsx | 255 ++++++++++++++++++ .../components/Query/TimeseriesQueryPage.tsx | 238 +++++++++++++--- .../main/resources/app/interfaces/types.d.ts | 52 +++- .../resources/app/utils/ChartConstants.ts | 45 ++++ .../resources/app/utils/TimeseriesUtils.ts | 161 +++++++++++ .../src/main/resources/package-lock.json | 104 ++++++- .../src/main/resources/package.json | 2 + 8 files changed, 1021 insertions(+), 48 deletions(-) create mode 100644 pinot-controller/src/main/resources/app/components/Query/MetricStatsTable.tsx create mode 100644 pinot-controller/src/main/resources/app/components/Query/TimeseriesChart.tsx create mode 100644 pinot-controller/src/main/resources/app/utils/ChartConstants.ts create mode 100644 pinot-controller/src/main/resources/app/utils/TimeseriesUtils.ts diff --git a/pinot-controller/src/main/resources/app/components/Query/MetricStatsTable.tsx b/pinot-controller/src/main/resources/app/components/Query/MetricStatsTable.tsx new file mode 100644 index 000000000000..0648950b2771 --- /dev/null +++ b/pinot-controller/src/main/resources/app/components/Query/MetricStatsTable.tsx @@ -0,0 +1,212 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import React from 'react'; +import { + Table, + TableBody, + TableCell, + TableContainer, + TableHead, + TableRow, + Paper, + Typography, + makeStyles, +} from '@material-ui/core'; + +import { ChartSeries } from 'Models'; +import { getSeriesColor } from '../../utils/ChartConstants'; + +const useStyles = makeStyles((theme) => ({ + tableContainer: { + marginTop: theme.spacing(2), + maxHeight: 400, + }, + table: { + minWidth: 650, + tableLayout: 'fixed', + }, + tableHead: { + backgroundColor: theme.palette.grey[100], + }, + tableHeadCell: { + color: theme.palette.text.primary, + fontSize: '0.875rem', + }, + tableRow: { + cursor: 'pointer', + '&:hover': { + backgroundColor: theme.palette.action.hover, + }, + '&.selected': { + backgroundColor: theme.palette.primary.light, + color: theme.palette.primary.contrastText, + }, + '&.dimmed': { + opacity: 0.4, + backgroundColor: theme.palette.grey[50], + }, + }, + metricCell: { + maxWidth: 300, + wordBreak: 'break-word', + fontSize: '0.875rem', + }, + statCell: { + fontSize: '0.875rem', + }, + colorBox: { + width: 10, + height: 10, + borderRadius: 2, + flexShrink: 0, + margin: '0 auto', + }, + + noDataMessage: { + textAlign: 'center', + padding: theme.spacing(3), + color: theme.palette.text.secondary, + }, + tableTitle: { + marginBottom: theme.spacing(1), + fontWeight: 'bold', + fontSize: '1.1rem', + }, +})); + +interface MetricStatsTableProps { + series: ChartSeries[]; + selectedMetric?: string; + onMetricSelect?: (metricName: string | null) => void; +} + +const MetricStatsTable: React.FC = ({ + series, + selectedMetric, + onMetricSelect +}) => { + const classes = useStyles(); + + const formatValue = (value: number): string => { + if (Math.abs(value) >= 1e6) { + return (value / 1e6).toFixed(2) + 'M'; + } else if (Math.abs(value) >= 1e3) { + return (value / 1e3).toFixed(2) + 'K'; + } else { + return value.toFixed(2); + } + }; + + + + const handleRowClick = (metricName: string) => { + if (onMetricSelect) { + // If clicking the same metric, deselect it + if (selectedMetric === metricName) { + onMetricSelect(null); + } else { + onMetricSelect(metricName); + } + } + }; + + const isRowSelected = (metricName: string) => selectedMetric === metricName; + const isRowDimmed = (metricName: string) => selectedMetric && selectedMetric !== metricName; + + if (!series || series.length === 0) { + return ( +
+ No metric data available + Run a timeseries query to see statistics +
+ ); + } + + return ( +
+ + + + + + Color + Metric + Count + Min + Max + Average + Sum + First Value + Last Value + + + + {series.map((chartSeries) => { + const { name, stats } = chartSeries; + const isSelected = isRowSelected(name); + const isDimmed = isRowDimmed(name); + + return ( + handleRowClick(name)} + > + +
s.name === name)) }} + /> + + + {name} + + + {stats.count} + + + {formatValue(stats.min)} + + + {formatValue(stats.max)} + + + {formatValue(stats.avg)} + + + {formatValue(stats.sum)} + + + {formatValue(stats.firstValue)} + + + {formatValue(stats.lastValue)} + + + ); + })} + +
+ +
+ ); +}; + +export default MetricStatsTable; diff --git a/pinot-controller/src/main/resources/app/components/Query/TimeseriesChart.tsx b/pinot-controller/src/main/resources/app/components/Query/TimeseriesChart.tsx new file mode 100644 index 000000000000..908f56c3e665 --- /dev/null +++ b/pinot-controller/src/main/resources/app/components/Query/TimeseriesChart.tsx @@ -0,0 +1,255 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import React, { useMemo } from 'react'; +import ReactECharts from 'echarts-for-react'; +import { makeStyles } from '@material-ui/core/styles'; +import { Typography, Paper } from '@material-ui/core'; +import { ChartSeries } from 'Models'; +import { getSeriesColor, MAX_SERIES_LIMIT, CHART_PADDING_PERCENTAGE } from '../../utils/ChartConstants'; + +// Define proper types for ECharts parameters +interface EChartsTooltipParams { + value: [number, number]; + seriesName: string; + color: string; +} + +// Extract chart configuration functions +const createChartSeries = (series: ChartSeries[], selectedMetric?: string) => { + const limitedSeries = series.slice(0, MAX_SERIES_LIMIT); + + return limitedSeries.map((s, index) => { + const isSelected = selectedMetric ? s.name === selectedMetric : true; + + return { + name: s.name, + type: 'line', + data: s.data.map(dp => [dp.timestamp, dp.value]), + smooth: false, + symbol: 'circle', + symbolSize: isSelected ? 8 : 4, + lineStyle: { + width: 1, + color: getSeriesColor(index), + opacity: isSelected ? 1 : 0, + }, + itemStyle: { + color: getSeriesColor(index), + opacity: isSelected ? 1 : 0, + }, + emphasis: { + focus: 'none', + scale: false, + }, + }; + }); +}; + +const createTooltipFormatter = (selectedMetric?: string) => { + return function (params: EChartsTooltipParams[]) { + let result = `
${new Date(params[0].value[0]).toLocaleString()}
`; + + // Filter params to only show selected series when one is selected + const filteredParams = selectedMetric + ? params.filter((param: EChartsTooltipParams) => param.seriesName === selectedMetric) + : params; + + filteredParams.forEach((param: EChartsTooltipParams) => { + const color = param.color; + const name = param.seriesName; + const value = param.value[1]; + result += `
+ + ${name}: ${value} +
`; + }); + return result; + }; +}; + +const getTimeRange = (series: ChartSeries[]) => { + const limitedSeries = series.slice(0, MAX_SERIES_LIMIT); + const allTimestamps = limitedSeries.flatMap(s => s.data.map(dp => dp.timestamp)); + const minTime = Math.min(...allTimestamps); + const maxTime = Math.max(...allTimestamps); + return { minTime, maxTime }; +}; + +const useStyles = makeStyles((theme) => ({ + chartContainer: { + height: '500px', + width: '100%', + padding: theme.spacing(2), + }, + noDataMessage: { + display: 'flex', + justifyContent: 'center', + alignItems: 'center', + height: '200px', + color: theme.palette.text.secondary, + }, + chartTitle: { + marginBottom: theme.spacing(2), + fontWeight: 'bold', + }, +})); + +interface TimeseriesChartProps { + series: ChartSeries[]; + height?: number; + selectedMetric?: string; +} + +const TimeseriesChart: React.FC = ({ + series, + height = 500, + selectedMetric +}) => { + const classes = useStyles(); + + const chartOption = useMemo(() => { + if (!series || series.length === 0) { + return {}; + } + + const chartSeries = createChartSeries(series, selectedMetric); + const { minTime, maxTime } = getTimeRange(series); + const timeRange = maxTime - minTime; + const padding = timeRange * CHART_PADDING_PERCENTAGE; + + return { + tooltip: { + trigger: 'axis', + formatter: createTooltipFormatter(selectedMetric), + axisPointer: { + type: 'cross', + label: { + backgroundColor: '#6a7985', + }, + }, + // Fix hover popup positioning to keep it within chart area + confine: true, + position: function (point: [number, number], params: unknown, dom: HTMLElement, rect: { x: number; y: number; width: number; height: number }, size: { viewSize: [number, number]; contentSize: [number, number] }) { + // Ensure tooltip stays within chart bounds + const [x, y] = point; + const [viewWidth, viewHeight] = size.viewSize; + const [contentWidth, contentHeight] = size.contentSize; + + let posX = x + 10; + let posY = y - 10; + + // Adjust horizontal position if tooltip would go outside + if (posX + contentWidth > viewWidth) { + posX = x - contentWidth - 10; + } + + // Adjust vertical position if tooltip would go outside + if (posY - contentHeight < 0) { + posY = y + 10; + } + + return [posX, posY]; + }, + }, + // Remove legend - colors will be shown in the stats table instead + grid: { + left: '3%', + right: '4%', + bottom: '3%', + top: '15%', + containLabel: true, + }, + xAxis: { + type: 'time', + boundaryGap: false, + min: minTime - padding, // Apply consistent padding + max: maxTime + padding, // Apply consistent padding + axisLabel: { + formatter: function (value: number) { + return new Date(value).toLocaleTimeString(); + }, + }, + }, + yAxis: { + type: 'value', + axisLabel: { + formatter: function (value: number) { + return value.toFixed(2); + }, + }, + }, + dataZoom: [ + { + type: 'inside', + start: 0, + end: 100, + zoomOnMouseWheel: true, + moveOnMouseMove: true, + moveOnMouseWheel: false, + preventDefaultMouseMove: false, + filterMode: 'filter', + rangeMode: ['value', 'value'], + }, + ], + toolbox: { + feature: { + dataZoom: { + yAxisIndex: 'none', + title: { + zoom: 'Zoom Selection' + } + }, + }, + right: 20, + top: 20, + }, + series: chartSeries, + }; + }, [series, selectedMetric]); + + if (!series || series.length === 0) { + return ( + +
+ No timeseries data available +
+ + ); + } + + const hasMoreSeries = series.length > MAX_SERIES_LIMIT; + + return ( + + {hasMoreSeries && ( + + Showing first {MAX_SERIES_LIMIT} of {series.length} series. Consider filtering your query to improve performance. + + )} + + + ); +}; + +export default TimeseriesChart; diff --git a/pinot-controller/src/main/resources/app/components/Query/TimeseriesQueryPage.tsx b/pinot-controller/src/main/resources/app/components/Query/TimeseriesQueryPage.tsx index c58db897c112..03922a0bc2eb 100644 --- a/pinot-controller/src/main/resources/app/components/Query/TimeseriesQueryPage.tsx +++ b/pinot-controller/src/main/resources/app/components/Query/TimeseriesQueryPage.tsx @@ -28,7 +28,9 @@ import { makeStyles, Button, Input, - FormControlLabel + FormControlLabel, + ButtonGroup, + Box } from '@material-ui/core'; import Alert from '@material-ui/lab/Alert'; import FileCopyIcon from '@material-ui/icons/FileCopy'; @@ -41,6 +43,41 @@ import { useHistory, useLocation } from 'react-router'; import TableToolbar from '../TableToolbar'; import { Resizable } from 're-resizable'; import SimpleAccordion from '../SimpleAccordion'; +import TimeseriesChart from './TimeseriesChart'; +import MetricStatsTable from './MetricStatsTable'; +import { parseTimeseriesResponse, isPrometheusFormat } from '../../utils/TimeseriesUtils'; +import { ChartSeries } from 'Models'; +import { MAX_SERIES_LIMIT } from '../../utils/ChartConstants'; + +// Define proper types +interface TimeseriesQueryResponse { + data: { + resultType: string; + result: Array<{ + metric: Record; + values: [number, number][]; + }>; + }; +} + +interface CodeMirrorEditor { + getValue: () => string; + setValue: (value: string) => void; +} + +interface CodeMirrorChangeData { + from: { line: number; ch: number }; + to: { line: number; ch: number }; + text: string[]; + removed: string[]; + origin: string; +} + +interface KeyboardEvent { + keyCode: number; + metaKey: boolean; + ctrlKey: boolean; +} const useStyles = makeStyles((theme) => ({ rightPanel: {}, @@ -106,8 +143,75 @@ const useStyles = makeStyles((theme) => ({ padding: theme.spacing(1), minWidth: 0, }, + })); +// Extract warning component +const TruncationWarning: React.FC<{ totalSeries: number; truncatedSeries: number }> = ({ + totalSeries, + truncatedSeries +}) => { + if (totalSeries <= truncatedSeries) return null; + + return ( + + + Large dataset detected: Showing first {truncatedSeries} of {totalSeries} series for visualization. + Switch to JSON view to see the complete dataset. + + + ); +}; + +// Extract view toggle component +const ViewToggle: React.FC<{ + viewType: 'json' | 'chart'; + onViewChange: (view: 'json' | 'chart') => void; + isChartDisabled: boolean; + onCopy: () => void; + copyMsg: boolean; + classes: ReturnType; +}> = ({ viewType, onViewChange, isChartDisabled, onCopy, copyMsg, classes }) => ( + + + + + + + + + + {copyMsg && ( + } + severity="info" + > + Copied results to Clipboard + + )} + + +); + const jsonoptions = { lineNumbers: true, mode: 'application/json', @@ -148,11 +252,16 @@ const TimeseriesQueryPage = () => { }); const [rawOutput, setRawOutput] = useState(''); - const [rawData, setRawData] = useState(null); + const [rawData, setRawData] = useState(null); + const [chartSeries, setChartSeries] = useState([]); + const [truncatedChartSeries, setTruncatedChartSeries] = useState([]); const [isLoading, setIsLoading] = useState(false); const [error, setError] = useState(''); const [shouldAutoExecute, setShouldAutoExecute] = useState(false); const [copyMsg, showCopyMsg] = React.useState(false); + const [viewType, setViewType] = useState<'json' | 'chart'>('chart'); + const [selectedMetric, setSelectedMetric] = useState(null); + // Update config when URL parameters change useEffect(() => { @@ -198,15 +307,15 @@ const TimeseriesQueryPage = () => { }); }, [history, location.pathname]); - const handleConfigChange = (field: keyof TimeseriesQueryConfig, value: any) => { + const handleConfigChange = (field: keyof TimeseriesQueryConfig, value: string | number | boolean) => { setConfig(prev => ({ ...prev, [field]: value })); }; - const handleQueryChange = (editor: any, data: any, value: string) => { + const handleQueryChange = (editor: CodeMirrorEditor, data: CodeMirrorChangeData, value: string) => { setConfig(prev => ({ ...prev, query: value })); }; - const handleQueryInterfaceKeyDown = (editor: any, event: any) => { + const handleQueryInterfaceKeyDown = (editor: CodeMirrorEditor, event: KeyboardEvent) => { const modifiedEnabled = event.metaKey == true || event.ctrlKey == true; // Map (Cmd/Ctrl) + Enter KeyPress to executing the query @@ -221,6 +330,25 @@ const TimeseriesQueryPage = () => { handleQueryInterfaceKeyDownRef.current = handleQueryInterfaceKeyDown; }, [handleQueryInterfaceKeyDown]); + // Extract data processing logic + const processQueryResponse = useCallback((parsedData: TimeseriesQueryResponse) => { + setRawData(parsedData); + setRawOutput(JSON.stringify(parsedData, null, 2)); + + // Parse timeseries data for chart and stats + if (isPrometheusFormat(parsedData)) { + const series = parseTimeseriesResponse(parsedData); + setChartSeries(series); + + // Create truncated series for visualization (limit to MAX_SERIES_LIMIT) + const truncatedSeries = series.slice(0, MAX_SERIES_LIMIT); + setTruncatedChartSeries(truncatedSeries); + } else { + setChartSeries([]); + setTruncatedChartSeries([]); + } + }, []); + const handleExecuteQuery = useCallback(async () => { if (!config.query.trim()) { setError('Please enter a query'); @@ -250,8 +378,7 @@ const TimeseriesQueryPage = () => { ? JSON.parse(response.data) : response.data; - setRawData(parsedData); - setRawOutput(JSON.stringify(parsedData, null, 2)); + processQueryResponse(parsedData); } catch (error) { console.error('Error executing timeseries query:', error); const errorMessage = error.response?.data?.message || error.message || 'Unknown error occurred'; @@ -259,7 +386,7 @@ const TimeseriesQueryPage = () => { } finally { setIsLoading(false); } - }, [config, updateURL]); + }, [config, updateURL, processQueryResponse]); const copyToClipboard = () => { const aux = document.createElement('input'); @@ -312,7 +439,7 @@ const TimeseriesQueryPage = () => { Query Language handleConfigChange('startTime', e.target.value)} + onChange={(e) => handleConfigChange('startTime', e.target.value as string)} placeholder={getOneMinuteAgoTimestamp()} /> @@ -341,7 +468,7 @@ const TimeseriesQueryPage = () => { handleConfigChange('endTime', e.target.value)} + onChange={(e) => handleConfigChange('endTime', e.target.value as string)} placeholder={getCurrentTimestamp()} /> @@ -353,7 +480,7 @@ const TimeseriesQueryPage = () => { handleConfigChange('timeout', parseInt(e.target.value) || 60000)} + onChange={(e) => handleConfigChange('timeout', parseInt(e.target.value as string) || 60000)} /> @@ -379,36 +506,65 @@ const TimeseriesQueryPage = () => { {rawOutput && ( - - - {copyMsg && ( - } - severity="info" - > - Copied results to Clipboard - - )} - - - - + {truncatedChartSeries.length > 0 ? ( + <> + + + + + ) : ( + + + No Chart Data Available + + + The query response is not in Prometheus-compatible format or contains no timeseries data. +
+ Switch to JSON view to see the raw response. + + + )} + + )} + + {viewType === 'json' && ( + + + + )} )} diff --git a/pinot-controller/src/main/resources/app/interfaces/types.d.ts b/pinot-controller/src/main/resources/app/interfaces/types.d.ts index 3808b38c04d3..e38d9d578b01 100644 --- a/pinot-controller/src/main/resources/app/interfaces/types.d.ts +++ b/pinot-controller/src/main/resources/app/interfaces/types.d.ts @@ -152,7 +152,7 @@ declare module 'Models' { serversUnparsableRespond: number; _segmentToConsumingInfoMap: Record; } - + /** * Pause status information for a realtime table */ @@ -225,6 +225,52 @@ declare module 'Models' { realtimeTotalCpuTimeNs: number }; + // Timeseries related types + export interface TimeseriesData { + metric: Record; + values: [number, number][]; // [timestamp, value] + } + + export interface TimeseriesResponse { + status: string; + data: { + resultType: string; + result: TimeseriesData[]; + }; + } + + export interface MetricStats { + min: number; + max: number; + avg: number; + sum: number; + count: number; + firstValue: number; + lastValue: number; + } + + export interface ChartDataPoint { + timestamp: number; + value: number; + formattedTime: string; + } + + export interface ChartSeries { + name: string; + data: ChartDataPoint[]; + stats: MetricStats; + } + + export type TimeseriesViewType = 'json' | 'chart' | 'stats'; + + export interface TimeseriesQueryConfig { + queryLanguage: string; + query: string; + startTime: string; + endTime: string; + timeout: number; + } + export type ClusterName = { clusterName: string }; @@ -313,7 +359,7 @@ declare module 'Models' { export type RebalanceTableSegmentJobs = { [key: string]: RebalanceTableSegmentJob; } - + export interface TaskRuntimeConfig { ConcurrentTasksPerWorker: string, TaskTimeoutMs: string, @@ -343,7 +389,7 @@ declare module 'Models' { OFFLINE = "OFFLINE", CONSUMING = "CONSUMING", ERROR = "ERROR" - } + } export const enum DISPLAY_SEGMENT_STATUS { BAD = "BAD", diff --git a/pinot-controller/src/main/resources/app/utils/ChartConstants.ts b/pinot-controller/src/main/resources/app/utils/ChartConstants.ts new file mode 100644 index 000000000000..5e1f823c851c --- /dev/null +++ b/pinot-controller/src/main/resources/app/utils/ChartConstants.ts @@ -0,0 +1,45 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +// Standard chart colors matching the ECharts palette +export const CHART_COLORS = [ + "#5470C6", "#91CC75", "#EE6666", "#FAC858", "#73C0DE", + "#3BA272", "#FC8452", "#9A60B4", "#EA7CCC", "#6E7074", + "#546570", "#C4CCD3", "#F05B72", "#FF715E", "#FFAF51", + "#FFE153", "#47B39C", "#5BACE1", "#32C5E9", "#96BFFF" +]; + +/** + * Maximum number of series that can be rendered in the chart + */ +export const MAX_SERIES_LIMIT = 20; + +/** + * Chart padding percentage for time axis and series + */ +export const CHART_PADDING_PERCENTAGE = 0.05; // 5% + +/** + * Get color for a series by index + * @param index - The series index + * @returns The color for the series + */ +export const getSeriesColor = (index: number): string => { + return CHART_COLORS[index % CHART_COLORS.length]; +}; diff --git a/pinot-controller/src/main/resources/app/utils/TimeseriesUtils.ts b/pinot-controller/src/main/resources/app/utils/TimeseriesUtils.ts new file mode 100644 index 000000000000..49dd041173b0 --- /dev/null +++ b/pinot-controller/src/main/resources/app/utils/TimeseriesUtils.ts @@ -0,0 +1,161 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import { ChartSeries, TimeseriesData, ChartDataPoint, MetricStats } from 'Models'; + +// Define proper types for API responses +interface PrometheusResponse { + data: { + resultType?: string; + result: TimeseriesData[]; + }; +} + +/** + * Parse Prometheus-compatible timeseries response + */ +export const parseTimeseriesResponse = (response: PrometheusResponse): ChartSeries[] => { + if (!response || !response.data || !response.data.result) { + return []; + } + + return response.data.result.map((series: TimeseriesData) => { + const dataPoints: ChartDataPoint[] = series.values.map(([timestamp, value]) => ({ + timestamp: timestamp * 1000, // Convert to milliseconds + value: parseFloat(String(value)), + formattedTime: new Date(timestamp * 1000).toLocaleString() + })); + + const stats = calculateMetricStats(dataPoints); + const seriesName = formatSeriesName(series.metric); + + return { + name: seriesName, + data: dataPoints, + stats + }; + }); +}; + +/** + * Calculate statistics for a metric series + */ +export const calculateMetricStats = (dataPoints: ChartDataPoint[]): MetricStats => { + if (dataPoints.length === 0) { + return { + min: NaN, + max: NaN, + avg: NaN, + sum: NaN, + count: 0, + firstValue: NaN, + lastValue: NaN + }; + } + + // Filter out null/undefined values and check if we have any valid values + const validValues = dataPoints + .map(dp => dp.value) + .filter(val => val !== null && val !== undefined && !isNaN(val)); + + if (validValues.length === 0) { + return { + min: NaN, + max: NaN, + avg: NaN, + sum: NaN, + count: 0, + firstValue: NaN, + lastValue: NaN + }; + } + + const sum = validValues.reduce((acc, val) => acc + val, 0); + const avg = sum / validValues.length; + const min = Math.min(...validValues); + const max = Math.max(...validValues); + + // Find first and last valid values + const firstValidIndex = dataPoints.findIndex(dp => dp.value !== null && dp.value !== undefined && !isNaN(dp.value)); + const lastValidIndex = dataPoints.length - 1 - dataPoints.slice().reverse().findIndex(dp => dp.value !== null && dp.value !== undefined && !isNaN(dp.value)); + + const firstValue = firstValidIndex >= 0 ? dataPoints[firstValidIndex].value : NaN; + const lastValue = lastValidIndex >= 0 ? dataPoints[lastValidIndex].value : NaN; + + return { + min, + max, + avg, + sum, + count: validValues.length, + firstValue, + lastValue + }; +}; + +/** + * Format series name from metric labels + */ +export const formatSeriesName = (metric: Record): string => { + if (!metric || Object.keys(metric).length === 0) { + return 'Unknown Metric'; + } + + // Try to find a meaningful name from common label patterns + const name = metric.__name__ || metric.name || metric.metric || metric.job || metric.instance; + + if (name) { + // Add additional context if available, but exclude the main metric name + const additionalLabels = Object.entries(metric) + .filter(([key]) => !['__name__', 'name', 'metric', 'job', 'instance'].includes(key)) + .map(([key, value]) => `${key}="${value}"`) + .join(', '); + + // Return just the name without the {} wrapper + return name; + } + + // Fallback to just the first label without {} + const firstLabel = Object.entries(metric)[0]; + return firstLabel ? firstLabel[0] : 'Unknown Metric'; +}; + +/** + * Check if response is in Prometheus format + */ +export const isPrometheusFormat = (response: PrometheusResponse): boolean => { + return response && + response.data && + Array.isArray(response.data.result); +}; + +/** + * Get time range from data points + */ +export const getTimeRange = (dataPoints: ChartDataPoint[]): { start: number; end: number } => { + if (dataPoints.length === 0) { + return { start: 0, end: 0 }; + } + + const timestamps = dataPoints.map(dp => dp.timestamp); + return { + start: Math.min(...timestamps), + end: Math.max(...timestamps) + }; +}; diff --git a/pinot-controller/src/main/resources/package-lock.json b/pinot-controller/src/main/resources/package-lock.json index 56e1d4ba3116..b38702a64c72 100644 --- a/pinot-controller/src/main/resources/package-lock.json +++ b/pinot-controller/src/main/resources/package-lock.json @@ -25,6 +25,8 @@ "codemirror": "5.65.5", "cross-fetch": "3.1.5", "dagre": "^0.8.5", + "echarts": "^5.4.3", + "echarts-for-react": "^3.0.2", "export-from-json": "1.6.0", "file": "0.2.2", "json-bigint": "1.0.0", @@ -3599,6 +3601,36 @@ "stream-shift": "^1.0.0" } }, + "node_modules/echarts": { + "version": "5.6.0", + "resolved": "https://registry.npmjs.org/echarts/-/echarts-5.6.0.tgz", + "integrity": "sha512-oTbVTsXfKuEhxftHqL5xprgLoc0k7uScAwtryCgWF6hPYFLRwOUHiFmHGCBKP5NPFNkDVopOieyUqYGH8Fa3kA==", + "license": "Apache-2.0", + "dependencies": { + "tslib": "2.3.0", + "zrender": "5.6.1" + } + }, + "node_modules/echarts-for-react": { + "version": "3.0.2", + "resolved": "https://registry.npmjs.org/echarts-for-react/-/echarts-for-react-3.0.2.tgz", + "integrity": "sha512-DRwIiTzx8JfwPOVgGttDytBqdp5VzCSyMRIxubgU/g2n9y3VLUmF2FK7Icmg/sNVkv4+rktmrLN9w22U2yy3fA==", + "license": "MIT", + "dependencies": { + "fast-deep-equal": "^3.1.3", + "size-sensor": "^1.0.1" + }, + "peerDependencies": { + "echarts": "^3.0.0 || ^4.0.0 || ^5.0.0", + "react": "^15.0.0 || >=16.0.0" + } + }, + "node_modules/echarts/node_modules/tslib": { + "version": "2.3.0", + "resolved": "https://registry.npmjs.org/tslib/-/tslib-2.3.0.tgz", + "integrity": "sha512-N82ooyxVNm6h1riLCoyS9e3fuJ3AMG2zIZs2Gd1ATcSFjSA23Q0fzjjZeh0jbJvWVDZ0cJT8yaNNaaXHzueNjg==", + "license": "0BSD" + }, "node_modules/ee-first": { "version": "1.1.1", "resolved": "https://registry.npmjs.org/ee-first/-/ee-first-1.1.1.tgz", @@ -4817,8 +4849,7 @@ "node_modules/fast-deep-equal": { "version": "3.1.3", "resolved": "https://registry.npmjs.org/fast-deep-equal/-/fast-deep-equal-3.1.3.tgz", - "integrity": "sha512-f3qQ9oQy9j2AhBe/H9VC91wLmKBCCU/gDOnKNAYG5hswO7BLKj09Hc5HYNz9cGI++xlpDCIgDaitVs03ATR84Q==", - "dev": true + "integrity": "sha512-f3qQ9oQy9j2AhBe/H9VC91wLmKBCCU/gDOnKNAYG5hswO7BLKj09Hc5HYNz9cGI++xlpDCIgDaitVs03ATR84Q==" }, "node_modules/fast-diff": { "version": "1.3.0", @@ -10694,6 +10725,12 @@ "integrity": "sha512-wnD2ZE+l+SPC/uoS0vXeE9L1+0wuaMqKlfz9AMUo38JsyLSBWSFcHR1Rri62LZc12vLr1gb3jl7iwQhgwpAbGQ==", "dev": true }, + "node_modules/size-sensor": { + "version": "1.0.2", + "resolved": "https://registry.npmjs.org/size-sensor/-/size-sensor-1.0.2.tgz", + "integrity": "sha512-2NCmWxY7A9pYKGXNBfteo4hy14gWu47rg5692peVMst6lQLPKrVjhY+UTEsPI5ceFRJSl3gVgMYaUi/hKuaiKw==", + "license": "ISC" + }, "node_modules/slash": { "version": "1.0.0", "resolved": "https://registry.npmjs.org/slash/-/slash-1.0.0.tgz", @@ -12956,6 +12993,21 @@ "node": ">=6" } }, + "node_modules/zrender": { + "version": "5.6.1", + "resolved": "https://registry.npmjs.org/zrender/-/zrender-5.6.1.tgz", + "integrity": "sha512-OFXkDJKcrlx5su2XbzJvj/34Q3m6PvyCZkVPHGYpcCJ52ek4U/ymZyfuV1nKE23AyBJ51E/6Yr0mhZ7xGTO4ag==", + "license": "BSD-3-Clause", + "dependencies": { + "tslib": "2.3.0" + } + }, + "node_modules/zrender/node_modules/tslib": { + "version": "2.3.0", + "resolved": "https://registry.npmjs.org/tslib/-/tslib-2.3.0.tgz", + "integrity": "sha512-N82ooyxVNm6h1riLCoyS9e3fuJ3AMG2zIZs2Gd1ATcSFjSA23Q0fzjjZeh0jbJvWVDZ0cJT8yaNNaaXHzueNjg==", + "license": "0BSD" + }, "node_modules/zustand": { "version": "3.7.2", "resolved": "https://registry.npmjs.org/zustand/-/zustand-3.7.2.tgz", @@ -15742,6 +15794,31 @@ "stream-shift": "^1.0.0" } }, + "echarts": { + "version": "5.6.0", + "resolved": "https://registry.npmjs.org/echarts/-/echarts-5.6.0.tgz", + "integrity": "sha512-oTbVTsXfKuEhxftHqL5xprgLoc0k7uScAwtryCgWF6hPYFLRwOUHiFmHGCBKP5NPFNkDVopOieyUqYGH8Fa3kA==", + "requires": { + "tslib": "2.3.0", + "zrender": "5.6.1" + }, + "dependencies": { + "tslib": { + "version": "2.3.0", + "resolved": "https://registry.npmjs.org/tslib/-/tslib-2.3.0.tgz", + "integrity": "sha512-N82ooyxVNm6h1riLCoyS9e3fuJ3AMG2zIZs2Gd1ATcSFjSA23Q0fzjjZeh0jbJvWVDZ0cJT8yaNNaaXHzueNjg==" + } + } + }, + "echarts-for-react": { + "version": "3.0.2", + "resolved": "https://registry.npmjs.org/echarts-for-react/-/echarts-for-react-3.0.2.tgz", + "integrity": "sha512-DRwIiTzx8JfwPOVgGttDytBqdp5VzCSyMRIxubgU/g2n9y3VLUmF2FK7Icmg/sNVkv4+rktmrLN9w22U2yy3fA==", + "requires": { + "fast-deep-equal": "^3.1.3", + "size-sensor": "^1.0.1" + } + }, "ee-first": { "version": "1.1.1", "resolved": "https://registry.npmjs.org/ee-first/-/ee-first-1.1.1.tgz", @@ -16671,8 +16748,7 @@ "fast-deep-equal": { "version": "3.1.3", "resolved": "https://registry.npmjs.org/fast-deep-equal/-/fast-deep-equal-3.1.3.tgz", - "integrity": "sha512-f3qQ9oQy9j2AhBe/H9VC91wLmKBCCU/gDOnKNAYG5hswO7BLKj09Hc5HYNz9cGI++xlpDCIgDaitVs03ATR84Q==", - "dev": true + "integrity": "sha512-f3qQ9oQy9j2AhBe/H9VC91wLmKBCCU/gDOnKNAYG5hswO7BLKj09Hc5HYNz9cGI++xlpDCIgDaitVs03ATR84Q==" }, "fast-diff": { "version": "1.3.0", @@ -21106,6 +21182,11 @@ "integrity": "sha512-wnD2ZE+l+SPC/uoS0vXeE9L1+0wuaMqKlfz9AMUo38JsyLSBWSFcHR1Rri62LZc12vLr1gb3jl7iwQhgwpAbGQ==", "dev": true }, + "size-sensor": { + "version": "1.0.2", + "resolved": "https://registry.npmjs.org/size-sensor/-/size-sensor-1.0.2.tgz", + "integrity": "sha512-2NCmWxY7A9pYKGXNBfteo4hy14gWu47rg5692peVMst6lQLPKrVjhY+UTEsPI5ceFRJSl3gVgMYaUi/hKuaiKw==" + }, "slash": { "version": "1.0.0", "resolved": "https://registry.npmjs.org/slash/-/slash-1.0.0.tgz", @@ -22784,6 +22865,21 @@ "decamelize": "^1.2.0" } }, + "zrender": { + "version": "5.6.1", + "resolved": "https://registry.npmjs.org/zrender/-/zrender-5.6.1.tgz", + "integrity": "sha512-OFXkDJKcrlx5su2XbzJvj/34Q3m6PvyCZkVPHGYpcCJ52ek4U/ymZyfuV1nKE23AyBJ51E/6Yr0mhZ7xGTO4ag==", + "requires": { + "tslib": "2.3.0" + }, + "dependencies": { + "tslib": { + "version": "2.3.0", + "resolved": "https://registry.npmjs.org/tslib/-/tslib-2.3.0.tgz", + "integrity": "sha512-N82ooyxVNm6h1riLCoyS9e3fuJ3AMG2zIZs2Gd1ATcSFjSA23Q0fzjjZeh0jbJvWVDZ0cJT8yaNNaaXHzueNjg==" + } + } + }, "zustand": { "version": "3.7.2", "resolved": "https://registry.npmjs.org/zustand/-/zustand-3.7.2.tgz", diff --git a/pinot-controller/src/main/resources/package.json b/pinot-controller/src/main/resources/package.json index d7ac3ff60eee..f1189f9a7445 100644 --- a/pinot-controller/src/main/resources/package.json +++ b/pinot-controller/src/main/resources/package.json @@ -80,6 +80,8 @@ "codemirror": "5.65.5", "cross-fetch": "3.1.5", "dagre": "^0.8.5", + "echarts": "^5.4.3", + "echarts-for-react": "^3.0.2", "export-from-json": "1.6.0", "file": "0.2.2", "json-bigint": "1.0.0", From af7047a23d233feffad0767c9220139acc66a07f Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 22 Jul 2025 16:50:23 -0700 Subject: [PATCH 002/167] Bump com.diffplug.spotless:spotless-maven-plugin from 2.46.0 to 2.46.1 (#16402) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 422626f6ff3f..2e7ade448e02 100644 --- a/pom.xml +++ b/pom.xml @@ -2156,7 +2156,7 @@ com.diffplug.spotless spotless-maven-plugin - 2.46.0 + 2.46.1 From b96200e89b30a74d787aaac785121c9ad0ad38ea Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 22 Jul 2025 16:50:40 -0700 Subject: [PATCH 003/167] Bump software.amazon.awssdk:bom from 2.32.4 to 2.32.5 (#16403) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 2e7ade448e02..a4ce61d7274a 100644 --- a/pom.xml +++ b/pom.xml @@ -177,7 +177,7 @@ 0.15.1 0.4.7 4.2.2 - 2.32.4 + 2.32.5 1.2.36 1.22.0 2.14.0 From b6e43ac60bbab3eabc6fa174b06a055bafc9eeb9 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 22 Jul 2025 16:51:07 -0700 Subject: [PATCH 004/167] Bump net.openhft:chronicle-bom from 2.27ea60 to 2.27ea62 (#16404) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index a4ce61d7274a..d21f59fb6e7c 100644 --- a/pom.xml +++ b/pom.xml @@ -295,7 +295,7 @@ 2.2.0 5.0.4 5.5.1 - 2.27ea60 + 2.27ea62 2.0.6.1 3.9.11 2.2.0 From b9e5116a4795be205c3b0ed9c3c8e9c8b484cef2 Mon Sep 17 00:00:00 2001 From: Hongkun Xu Date: Wed, 23 Jul 2025 08:09:41 +0800 Subject: [PATCH 005/167] serialize TlsSpec when executing DataIngestion backfill job (#16381) --- .../org/apache/pinot/spi/ingestion/batch/spec/TlsSpec.java | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/ingestion/batch/spec/TlsSpec.java b/pinot-spi/src/main/java/org/apache/pinot/spi/ingestion/batch/spec/TlsSpec.java index 4c446f56baee..bcc1f890206c 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/ingestion/batch/spec/TlsSpec.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/ingestion/batch/spec/TlsSpec.java @@ -18,11 +18,14 @@ */ package org.apache.pinot.spi.ingestion.batch.spec; +import java.io.Serializable; + + /** * TLS key and trust-store specification for ingestion jobs * (Enables access to TLS-protected controller APIs, etc.) */ -public class TlsSpec { +public class TlsSpec implements Serializable { private String _keyStorePath; private String _keyStorePassword; private String _trustStoreType; From 831a25aff4758deddaadf0631edfacfeef5b0cfb Mon Sep 17 00:00:00 2001 From: NOOB <43700604+noob-se7en@users.noreply.github.com> Date: Wed, 23 Jul 2025 06:16:47 +0530 Subject: [PATCH 006/167] Fixes race condition in RealtimeConsumptionRateManager.MetricEmitter used by serverRateLimiter (#16339) --- .../RealtimeConsumptionRateManager.java | 274 +++++++++++++++--- .../RealtimeConsumptionRateManagerTest.java | 91 +++++- ...rverRateLimitConfigChangeListenerTest.java | 68 ++++- 3 files changed, 365 insertions(+), 68 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeConsumptionRateManager.java b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeConsumptionRateManager.java index aeaaebcc35be..95604153a688 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeConsumptionRateManager.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeConsumptionRateManager.java @@ -22,13 +22,18 @@ import com.google.common.cache.CacheBuilder; import com.google.common.cache.CacheLoader; import com.google.common.cache.LoadingCache; +import com.google.common.util.concurrent.AtomicDouble; import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.RateLimiter; import java.time.Clock; import java.time.Instant; import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.LongAdder; import org.apache.pinot.common.metrics.ServerGauge; import org.apache.pinot.common.metrics.ServerMeter; import org.apache.pinot.common.metrics.ServerMetrics; @@ -87,24 +92,35 @@ public ConsumptionRateLimiter createServerRateLimiter(PinotConfiguration serverC double serverRateLimit = serverConfig.getProperty(CommonConstants.Server.CONFIG_OF_SERVER_CONSUMPTION_RATE_LIMIT, CommonConstants.Server.DEFAULT_SERVER_CONSUMPTION_RATE_LIMIT); - _serverRateLimiter = createServerRateLimiter(serverRateLimit, serverMetrics); + createOrUpdateServerRateLimiter(serverRateLimit, serverMetrics); return _serverRateLimiter; } - private ConsumptionRateLimiter createServerRateLimiter(double serverRateLimit, ServerMetrics serverMetrics) { + private void createOrUpdateServerRateLimiter(double serverRateLimit, ServerMetrics serverMetrics) { + LOGGER.info("Setting up ServerRateLimiter with rate limit: {}", serverRateLimit); + ConsumptionRateLimiter currentRateLimiter = _serverRateLimiter; if (serverRateLimit > 0) { - LOGGER.info("Set up ConsumptionRateLimiter with rate limit: {}", serverRateLimit); - MetricEmitter metricEmitter = new MetricEmitter(serverMetrics, SERVER_CONSUMPTION_RATE_METRIC_KEY_NAME); - return new RateLimiterImpl(serverRateLimit, metricEmitter); + if (currentRateLimiter instanceof ServerRateLimiter) { + ((ServerRateLimiter) currentRateLimiter).updateRateLimit(serverRateLimit); + } else { + _serverRateLimiter = + new ServerRateLimiter(serverRateLimit, serverMetrics, SERVER_CONSUMPTION_RATE_METRIC_KEY_NAME); + } } else { - LOGGER.info("ConsumptionRateLimiter is disabled"); - return NOOP_RATE_LIMITER; + if (currentRateLimiter instanceof ServerRateLimiter) { + ((ServerRateLimiter) currentRateLimiter).close(); + _serverRateLimiter = NOOP_RATE_LIMITER; + // Note: The consumer threads already present before refer to the serverRateLimiter object (not + // NOOP_RATE_LIMITER). These threads will keep calling throttle() method and they might get blocked as a + // result of it, But the metric related to throttling won't be emitted since as a result of here above the + // AsyncMetricEmitter will be closed. It's recommended to forceCommit segments to avoid this. + } } } - public void updateServerRateLimiter(double newServerRateLimit, ServerMetrics serverMetrics) { - LOGGER.info("Updating serverRateLimiter from: {} to: {}", _serverRateLimiter, newServerRateLimit); - _serverRateLimiter = createServerRateLimiter(newServerRateLimit, serverMetrics); + public void updateServerRateLimiter(double newRateLimit, ServerMetrics serverMetrics) { + LOGGER.info("Updating serverRateLimiter from: {} to: {}", _serverRateLimiter, newRateLimit); + createOrUpdateServerRateLimiter(newRateLimit, serverMetrics); } public ConsumptionRateLimiter getServerRateLimiter() { @@ -128,8 +144,7 @@ public ConsumptionRateLimiter createRateLimiter(StreamConfig streamConfig, Strin LOGGER.info("A consumption rate limiter is set up for topic {} in table {} with rate limit: {} " + "(topic rate limit: {}, partition count: {})", streamConfig.getTopicName(), tableName, partitionRateLimit, topicRateLimit, partitionCount); - MetricEmitter metricEmitter = new MetricEmitter(serverMetrics, metricKeyName); - return new RateLimiterImpl(partitionRateLimit, metricEmitter); + return new PartitionRateLimiter(partitionRateLimit, serverMetrics, metricKeyName); } @VisibleForTesting @@ -162,6 +177,58 @@ public ListenableFuture reload(StreamConfig key, Integer oldValue) }); } + /** + * Tracks quota utilization for a stream consumer by aggregating the number of messages consumed and computing + * the consumption rate against the configured rate limit. + *

+ * This class maintains the message count over time and computes the quota utilization ratio once per minute. + * The utilization is reported as a gauge metric to {@link ServerMetrics}. + *

+ * Note: This class is not thread-safe and is intended to be used in contexts where concurrency control is + * managed externally (e.g., per-partition rate limiter instances). + *

+ * Example: + * - If 3000 messages are consumed in one minute and the rate limit is 50 msg/sec, + * the quota utilization = (3000 / 60) / 50 = 1.0 → 100% + */ + static class QuotaUtilizationTracker { + private long _previousMinute = -1; + private int _aggregateNumMessages = 0; + private final ServerMetrics _serverMetrics; + private final String _metricKeyName; + + public QuotaUtilizationTracker(ServerMetrics serverMetrics, String metricKeyName) { + _serverMetrics = serverMetrics; + _metricKeyName = metricKeyName; + } + + /** + * Update count and return utilization ratio percentage (0 if not enough data yet). + */ + public int update(int numMsgsConsumed, double rateLimit, Instant now) { + int ratioPercentage = 0; + long nowInMinutes = now.getEpochSecond() / 60; + if (nowInMinutes == _previousMinute) { + _aggregateNumMessages += numMsgsConsumed; + } else { + if (_previousMinute != -1) { // not first time + double actualRate = _aggregateNumMessages / ((nowInMinutes - _previousMinute) * 60.0); // messages per second + ratioPercentage = (int) Math.round(actualRate / rateLimit * 100); + _serverMetrics.setValueOfTableGauge(_metricKeyName, ServerGauge.CONSUMPTION_QUOTA_UTILIZATION, + ratioPercentage); + } + _aggregateNumMessages = numMsgsConsumed; + _previousMinute = nowInMinutes; + } + return ratioPercentage; + } + + @VisibleForTesting + int getAggregateNumMessages() { + return _aggregateNumMessages; + } + } + @FunctionalInterface public interface ConsumptionRateLimiter { void throttle(int numMsgs); @@ -171,16 +238,23 @@ public interface ConsumptionRateLimiter { static final ConsumptionRateLimiter NOOP_RATE_LIMITER = n -> { }; + /** + * {@code PartitionRateLimiter} is an implementation of {@link ConsumptionRateLimiter} that uses Guava's + * {@link com.google.common.util.concurrent.RateLimiter} to throttle the rate of consumption at per partition + * level based on a configurable rate limit (in permits per second). + *

+ *

This class is NOT thread-safe + */ @VisibleForTesting - static class RateLimiterImpl implements ConsumptionRateLimiter { + static class PartitionRateLimiter implements ConsumptionRateLimiter { private final double _rate; private final RateLimiter _rateLimiter; - private MetricEmitter _metricEmitter; + private final QuotaUtilizationTracker _quotaUtilizationTracker; - private RateLimiterImpl(double rate, MetricEmitter metricEmitter) { + private PartitionRateLimiter(double rate, ServerMetrics serverMetrics, String metricKeyName) { _rate = rate; _rateLimiter = RateLimiter.create(rate); - _metricEmitter = metricEmitter; + _quotaUtilizationTracker = new QuotaUtilizationTracker(serverMetrics, metricKeyName); } @Override @@ -189,7 +263,7 @@ public void throttle(int numMsgs) { // Only emit metrics when throttling is allowed. Throttling is not enabled. // until the server has passed startup checks. Otherwise, we will see // consumption well over 100% during startup. - _metricEmitter.emitMetric(numMsgs, _rate, Clock.systemUTC().instant()); + _quotaUtilizationTracker.update(numMsgs, _rate, Clock.systemUTC().instant()); if (numMsgs > 0) { _rateLimiter.acquire(numMsgs); } @@ -203,9 +277,67 @@ public void throttle(int numMsgs) { @Override public String toString() { - return "RateLimiterImpl{" + return "PartitionRateLimiter{" + "_rate=" + _rate + ", _rateLimiter=" + _rateLimiter + + ", _quotaUtilizationTracker=" + _quotaUtilizationTracker + + '}'; + } + } + + /** + * {@code ServerRateLimiter} is an implementation of {@link ConsumptionRateLimiter} that uses Guava's + * {@link com.google.common.util.concurrent.RateLimiter} to throttle the rate of consumption at entire server + * level based on a configurable rate limit (in permits per second). + *

+ * It supports dynamically updating the rate limit and emits metrics asynchronously to track quota utilization + * via {@link AsyncMetricEmitter}. + * + *

This class is thread-safe + */ + @VisibleForTesting + static class ServerRateLimiter implements ConsumptionRateLimiter { + private final RateLimiter _rateLimiter; + private final AsyncMetricEmitter _metricEmitter; + + public ServerRateLimiter(double initialRateLimit, ServerMetrics serverMetrics, String metricKeyName) { + _rateLimiter = RateLimiter.create(initialRateLimit); + _metricEmitter = new AsyncMetricEmitter(serverMetrics, metricKeyName, initialRateLimit); + _metricEmitter.start(); // start background emission + } + + public void throttle(int numMsgsConsumed) { + if (InstanceHolder.INSTANCE._isThrottlingAllowed) { + _metricEmitter.record(numMsgsConsumed); // just incrementing counter (non-blocking) + if (numMsgsConsumed > 0) { + _rateLimiter.acquire(numMsgsConsumed); // blocks if needed + } + } + } + + public void updateRateLimit(double newRateLimit) { + _rateLimiter.setRate(newRateLimit); + _metricEmitter.updateRateLimit(newRateLimit); + } + + public void close() { + _metricEmitter.close(); + } + + @VisibleForTesting + double getRate() { + return _rateLimiter.getRate(); + } + + @VisibleForTesting + AsyncMetricEmitter getMetricEmitter() { + return _metricEmitter; + } + + @Override + public String toString() { + return "ServerRateLimiter{" + + "_rateLimiter=" + _rateLimiter + ", _metricEmitter=" + _metricEmitter + '}'; } @@ -231,41 +363,91 @@ interface PartitionCountFetcher { }; /** - * This class is responsible to emit a gauge metric for the ratio of the actual consumption rate to the rate limit. - * Number of messages consumed are aggregated over one minute. Each minute the ratio percentage is calculated and - * emitted. + * Asynchronously emits the quota utilization metric for a shared rate limiter (e.g., server-wide). + *

+ * This class aggregates consumed message counts over a fixed time interval (default: 60 seconds) using a + * high-performance {@link java.util.concurrent.atomic.LongAdder}. A scheduled background task computes the + * actual message rate and reports the quota utilization ratio as a gauge metric. + *

+ * This design avoids contention from multiple threads calling emit logic and ensures non-blocking accumulation + * of message counts. Thread-safe and suitable for shared use. + *

+ * Usage: + * - Call {@link #record(int)} to record messages consumed. + * - Call {@link #start()} once to schedule metric emission. + * - Optionally call {@link #close()} to stop the emitter. + *

+ * Thread-safe. */ - @VisibleForTesting - static class MetricEmitter { + static class AsyncMetricEmitter { + private static final int METRIC_EMIT_FREQUENCY_SEC = 60; + private final AtomicDouble _rateLimit; + private final LongAdder _messageCount = new LongAdder(); + private final ScheduledExecutorService _executor; + private final AtomicBoolean _running = new AtomicBoolean(false); + private final QuotaUtilizationTracker _tracker; - private final ServerMetrics _serverMetrics; - private final String _metricKeyName; + public AsyncMetricEmitter(ServerMetrics serverMetrics, String metricKeyName, double initialRateLimit) { + _rateLimit = new AtomicDouble(initialRateLimit); + _tracker = new QuotaUtilizationTracker(serverMetrics, metricKeyName); + _executor = Executors.newSingleThreadScheduledExecutor(r -> { + Thread t = new Thread(r, "server-rate-limit-metric-emitter"); + t.setDaemon(true); + return t; + }); + } - // state variables - private long _previousMinute = -1; - private int _aggregateNumMessages = 0; + public void start() { + if (_running.compareAndSet(false, true)) { + _executor.scheduleAtFixedRate(this::emit, 0, METRIC_EMIT_FREQUENCY_SEC, TimeUnit.SECONDS); + } + } - public MetricEmitter(ServerMetrics serverMetrics, String metricKeyName) { - _serverMetrics = serverMetrics; - _metricKeyName = metricKeyName; + @VisibleForTesting + void start(int initialDelayInSeconds, int emitFrequencyInSeconds) { + if (_running.compareAndSet(false, true)) { + _executor.scheduleAtFixedRate(this::emit, initialDelayInSeconds, emitFrequencyInSeconds, TimeUnit.SECONDS); + } } - int emitMetric(int numMsgsConsumed, double rateLimit, Instant now) { - int ratioPercentage = 0; - long nowInMinutes = now.getEpochSecond() / 60; - if (nowInMinutes == _previousMinute) { - _aggregateNumMessages += numMsgsConsumed; - } else { - if (_previousMinute != -1) { // not first time - double actualRate = _aggregateNumMessages / ((nowInMinutes - _previousMinute) * 60.0); // messages per second - ratioPercentage = (int) Math.round(actualRate / rateLimit * 100); - _serverMetrics.setValueOfTableGauge(_metricKeyName, ServerGauge.CONSUMPTION_QUOTA_UTILIZATION, - ratioPercentage); - } - _aggregateNumMessages = numMsgsConsumed; - _previousMinute = nowInMinutes; + public void updateRateLimit(double newRateLimit) { + _rateLimit.set(newRateLimit); + } + + public void record(int numMsgsConsumed) { + _messageCount.add(numMsgsConsumed); + } + + private void emit() { + try { + double rateLimit = _rateLimit.get(); + Instant now = Instant.now(); + int count = (int) _messageCount.sumThenReset(); + _tracker.update(count, rateLimit, now); + } catch (Exception e) { + LOGGER.warn("Encountered an error while emitting the metrics.", e); + } + } + + @VisibleForTesting + LongAdder getMessageCount() { + return _messageCount; + } + + @VisibleForTesting + QuotaUtilizationTracker getTracker() { + return _tracker; + } + + @VisibleForTesting + double getRate() { + return _rateLimit.get(); + } + + public void close() { + if (_running.compareAndSet(true, false)) { + _executor.shutdownNow(); } - return ratioPercentage; } } } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/RealtimeConsumptionRateManagerTest.java b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/RealtimeConsumptionRateManagerTest.java index 325066fa1622..8face43359fc 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/RealtimeConsumptionRateManagerTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/RealtimeConsumptionRateManagerTest.java @@ -24,12 +24,14 @@ import java.time.ZoneOffset; import java.util.Arrays; import java.util.Optional; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; import org.apache.pinot.common.metrics.ServerMetrics; import org.apache.pinot.spi.env.PinotConfiguration; import org.apache.pinot.spi.stream.StreamConfig; import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.util.TestUtils; import org.testng.annotations.Test; import static org.apache.pinot.core.data.manager.realtime.RealtimeConsumptionRateManager.*; @@ -82,11 +84,11 @@ public class RealtimeConsumptionRateManagerTest { public void testCreateRateLimiter() { // topic A ConsumptionRateLimiter rateLimiter = _consumptionRateManager.createRateLimiter(STREAM_CONFIG_A, TABLE_NAME); - assertEquals(5.0, ((RateLimiterImpl) rateLimiter).getRate(), DELTA); + assertEquals(5.0, ((PartitionRateLimiter) rateLimiter).getRate(), DELTA); // topic B rateLimiter = _consumptionRateManager.createRateLimiter(STREAM_CONFIG_B, TABLE_NAME); - assertEquals(2.5, ((RateLimiterImpl) rateLimiter).getRate(), DELTA); + assertEquals(2.5, ((PartitionRateLimiter) rateLimiter).getRate(), DELTA); // topic C rateLimiter = _consumptionRateManager.createRateLimiter(STREAM_CONFIG_C, TABLE_NAME); @@ -97,11 +99,21 @@ public void testCreateRateLimiter() { public void testCreateServerRateLimiter() { // Server config 1 ConsumptionRateLimiter rateLimiter = _consumptionRateManager.createServerRateLimiter(SERVER_CONFIG_1, null); - assertEquals(5.0, ((RateLimiterImpl) rateLimiter).getRate(), DELTA); + ServerRateLimiter serverRateLimiter = (ServerRateLimiter) rateLimiter; + try { + assertEquals(serverRateLimiter.getRate(), 5.0, DELTA); + assertEquals(serverRateLimiter.getMetricEmitter().getRate(), 5.0, DELTA); + } finally { + serverRateLimiter.close(); + } // Server config 2 - rateLimiter = _consumptionRateManager.createServerRateLimiter(SERVER_CONFIG_2, null); - assertEquals(2.5, ((RateLimiterImpl) rateLimiter).getRate(), DELTA); + serverRateLimiter = (ServerRateLimiter) _consumptionRateManager.createServerRateLimiter(SERVER_CONFIG_2, null); + try { + assertEquals(((ServerRateLimiter) rateLimiter).getRate(), 2.5, DELTA); + } finally { + serverRateLimiter.close(); + } // Server config 3 rateLimiter = _consumptionRateManager.createServerRateLimiter(SERVER_CONFIG_3, null); @@ -110,6 +122,19 @@ public void testCreateServerRateLimiter() { // Server config 4 rateLimiter = _consumptionRateManager.createServerRateLimiter(SERVER_CONFIG_4, null); assertEquals(rateLimiter, NOOP_RATE_LIMITER); + + _consumptionRateManager.updateServerRateLimiter(1, null); + serverRateLimiter = (ServerRateLimiter) _consumptionRateManager.getServerRateLimiter(); + try { + assertEquals(serverRateLimiter.getRate(), 1); + assertEquals(serverRateLimiter.getMetricEmitter().getRate(), 1); + + serverRateLimiter.updateRateLimit(12_000); + assertEquals(serverRateLimiter.getRate(), 12_000); + assertEquals(serverRateLimiter.getMetricEmitter().getRate(), 12_000); + } finally { + serverRateLimiter.close(); + } } @Test @@ -164,48 +189,84 @@ public void testMetricEmitter() { double rateLimit = 2; // unit: msgs/sec double rateLimitInMinutes = rateLimit * 60; ServerMetrics serverMetrics = mock(ServerMetrics.class); - MetricEmitter metricEmitter = new MetricEmitter(serverMetrics, "tableA-topicB-partition5"); + QuotaUtilizationTracker quotaUtilizationTracker = + new QuotaUtilizationTracker(serverMetrics, "tableA-topicB-partition5"); // 1st minute: no metrics should be emitted in the first minute int[] numMsgs = {10, 20, 5, 25}; Instant now = Clock.fixed(Instant.parse("2022-08-10T12:00:02Z"), ZoneOffset.UTC).instant(); - assertEquals(metricEmitter.emitMetric(numMsgs[0], rateLimit, now), 0); + assertEquals(quotaUtilizationTracker.update(numMsgs[0], rateLimit, now), 0); now = Clock.fixed(Instant.parse("2022-08-10T12:00:10Z"), ZoneOffset.UTC).instant(); - assertEquals(metricEmitter.emitMetric(numMsgs[1], rateLimit, now), 0); + assertEquals(quotaUtilizationTracker.update(numMsgs[1], rateLimit, now), 0); now = Clock.fixed(Instant.parse("2022-08-10T12:00:30Z"), ZoneOffset.UTC).instant(); - assertEquals(metricEmitter.emitMetric(numMsgs[2], rateLimit, now), 0); + assertEquals(quotaUtilizationTracker.update(numMsgs[2], rateLimit, now), 0); now = Clock.fixed(Instant.parse("2022-08-10T12:00:55Z"), ZoneOffset.UTC).instant(); - assertEquals(metricEmitter.emitMetric(numMsgs[3], rateLimit, now), 0); + assertEquals(quotaUtilizationTracker.update(numMsgs[3], rateLimit, now), 0); // 2nd minute: metric should be emitted now = Clock.fixed(Instant.parse("2022-08-10T12:01:05Z"), ZoneOffset.UTC).instant(); int sumOfMsgsInPrevMinute = sum(numMsgs); int expectedRatio = calcExpectedRatio(rateLimitInMinutes, sumOfMsgsInPrevMinute); numMsgs = new int[]{35}; - assertEquals(metricEmitter.emitMetric(numMsgs[0], rateLimit, now), expectedRatio); + assertEquals(quotaUtilizationTracker.update(numMsgs[0], rateLimit, now), expectedRatio); // 3rd minute now = Clock.fixed(Instant.parse("2022-08-10T12:02:25Z"), ZoneOffset.UTC).instant(); sumOfMsgsInPrevMinute = sum(numMsgs); expectedRatio = calcExpectedRatio(rateLimitInMinutes, sumOfMsgsInPrevMinute); numMsgs = new int[]{0}; - assertEquals(metricEmitter.emitMetric(numMsgs[0], rateLimit, now), expectedRatio); + assertEquals(quotaUtilizationTracker.update(numMsgs[0], rateLimit, now), expectedRatio); // 4th minute now = Clock.fixed(Instant.parse("2022-08-10T12:03:15Z"), ZoneOffset.UTC).instant(); sumOfMsgsInPrevMinute = sum(numMsgs); expectedRatio = calcExpectedRatio(rateLimitInMinutes, sumOfMsgsInPrevMinute); numMsgs = new int[]{10, 20}; - assertEquals(metricEmitter.emitMetric(numMsgs[0], rateLimit, now), expectedRatio); + assertEquals(quotaUtilizationTracker.update(numMsgs[0], rateLimit, now), expectedRatio); now = Clock.fixed(Instant.parse("2022-08-10T12:03:20Z"), ZoneOffset.UTC).instant(); - assertEquals(metricEmitter.emitMetric(numMsgs[1], rateLimit, now), expectedRatio); + assertEquals(quotaUtilizationTracker.update(numMsgs[1], rateLimit, now), expectedRatio); // 5th minute now = Clock.fixed(Instant.parse("2022-08-10T12:04:30Z"), ZoneOffset.UTC).instant(); sumOfMsgsInPrevMinute = sum(numMsgs); expectedRatio = calcExpectedRatio(rateLimitInMinutes, sumOfMsgsInPrevMinute); numMsgs = new int[]{5}; - assertEquals(metricEmitter.emitMetric(numMsgs[0], rateLimit, now), expectedRatio); + assertEquals(quotaUtilizationTracker.update(numMsgs[0], rateLimit, now), expectedRatio); + } + + @Test + public void testAsyncMetricEmitter() + throws InterruptedException { + AsyncMetricEmitter emitter = new AsyncMetricEmitter(mock(ServerMetrics.class), "testMetric", 10.0); + try { + emitter.start(0, 1); + Thread.sleep(1500); // Let emitter run at-least once + for (int i = 0; i < 20; i++) { + CompletableFuture.runAsync(() -> emitter.record(1)); + } + TestUtils.waitForCondition( + aVoid -> (emitter.getMessageCount().intValue() == 0) && (emitter.getTracker().getAggregateNumMessages() > 0), + 5000, + "Expected messageCount to be zero because messageCount is always reset before emitter calls " + + "quotaUtilisationTracker"); + } finally { + emitter.close(); + } + + AsyncMetricEmitter emitter1 = new AsyncMetricEmitter(mock(ServerMetrics.class), "testMetric", 10.0); + try { + emitter1.start(10, 10); + for (int i = 0; i < 20; i++) { + CompletableFuture.runAsync(() -> emitter1.record(1)); + } + TestUtils.waitForCondition( + aVoid -> ((emitter1.getMessageCount().intValue() > 0) && (emitter1.getTracker().getAggregateNumMessages() + == 0)), 5000, + "Expected messageCount to be greater than zero because messageCount will reset post initial delay (first " + + "run)."); + } finally { + emitter.close(); + } } private int calcExpectedRatio(double rateLimitInMinutes, int sumOfMsgsInPrevMinute) { diff --git a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/ServerRateLimitConfigChangeListenerTest.java b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/ServerRateLimitConfigChangeListenerTest.java index e2a2d2c5df9a..192bf3283b23 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/ServerRateLimitConfigChangeListenerTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/ServerRateLimitConfigChangeListenerTest.java @@ -23,6 +23,8 @@ import java.util.List; import java.util.Map; import java.util.Set; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicReference; import org.apache.pinot.common.metrics.ServerMetrics; import org.apache.pinot.spi.env.PinotConfiguration; import org.apache.pinot.spi.utils.CommonConstants; @@ -45,10 +47,13 @@ public class ServerRateLimitConfigChangeListenerTest { } @Test - public void testRateLimitUpdate() { + public void testRateLimitUpdate() + throws InterruptedException { + AtomicReference errorRef = new AtomicReference<>(); + simulateThrottling(errorRef); // Initial state RealtimeConsumptionRateManager.getInstance().createServerRateLimiter(SERVER_CONFIG, null); - RealtimeConsumptionRateManager.RateLimiterImpl serverRateLimiter = getServerRateLimiter(); + RealtimeConsumptionRateManager.ServerRateLimiter serverRateLimiter = getServerRateLimiter(); double initialRate = serverRateLimiter.getRate(); assertEquals(initialRate, 5.0, DELTA); @@ -58,18 +63,67 @@ public void testRateLimitUpdate() { ServerRateLimitConfigChangeListener listener = new ServerRateLimitConfigChangeListener(MOCK_SERVER_METRICS); Set changedConfigSet = new HashSet<>(List.of(CommonConstants.Server.CONFIG_OF_SERVER_CONSUMPTION_RATE_LIMIT)); + simulateThrottling(errorRef); listener.onChange(changedConfigSet, newConfig); + simulateThrottling(errorRef); - // Verify that old rate remains same and the new rate is applied + // Verify that rate changed double rate = serverRateLimiter.getRate(); - assertEquals(rate, 5.0, DELTA); - + assertEquals(rate, 300.0, DELTA); double updatedRate = getServerRateLimiter().getRate(); assertEquals(updatedRate, 300.0, DELTA); + + // Test removal of serverRateLimit + newConfig = new HashMap<>(); + newConfig.put(CommonConstants.Server.CONFIG_OF_SERVER_CONSUMPTION_RATE_LIMIT, "0"); + changedConfigSet = new HashSet<>(List.of(CommonConstants.Server.CONFIG_OF_SERVER_CONSUMPTION_RATE_LIMIT)); + simulateThrottling(errorRef); + listener.onChange(changedConfigSet, newConfig); + simulateThrottling(errorRef); + + // Verify that old rate remains same and the new rate is applied + rate = serverRateLimiter.getRate(); + assertEquals(rate, 300.0, DELTA); + + assertEquals(RealtimeConsumptionRateManager.NOOP_RATE_LIMITER, + RealtimeConsumptionRateManager.getInstance().getServerRateLimiter()); + + // Test update of serverRateLimit after it was removed + newConfig = new HashMap<>(); + newConfig.put(CommonConstants.Server.CONFIG_OF_SERVER_CONSUMPTION_RATE_LIMIT, "10000"); + changedConfigSet = new HashSet<>(List.of(CommonConstants.Server.CONFIG_OF_SERVER_CONSUMPTION_RATE_LIMIT)); + simulateThrottling(errorRef); + listener.onChange(changedConfigSet, newConfig); + simulateThrottling(errorRef); + + // Verify that old rate (one before the config change was deleted and again added) remains same. + rate = serverRateLimiter.getRate(); + assertEquals(rate, 300.0, DELTA); + + updatedRate = getServerRateLimiter().getRate(); + assertEquals(updatedRate, 10000, DELTA); + + Thread.sleep(1000); + if (errorRef.get() != null) { + throw new RuntimeException("Throttle call failed: " + errorRef.get().getMessage()); + } + } + + private void simulateThrottling(AtomicReference errorRef) { + // A helper method to test side effects of throttling during serverRateLimit config change. + for (int i = 0; i < 10; i++) { + CompletableFuture.runAsync(() -> { + try { + RealtimeConsumptionRateManager.getInstance().getServerRateLimiter().throttle(100); + } catch (Throwable throwable) { + errorRef.set(throwable); + } + }); + } } - private RealtimeConsumptionRateManager.RateLimiterImpl getServerRateLimiter() { - return (RealtimeConsumptionRateManager.RateLimiterImpl) (RealtimeConsumptionRateManager.getInstance() + private RealtimeConsumptionRateManager.ServerRateLimiter getServerRateLimiter() { + return (RealtimeConsumptionRateManager.ServerRateLimiter) (RealtimeConsumptionRateManager.getInstance() .getServerRateLimiter()); } } From 90346d5f379d24687b0117ee7ff8aafa113ac552 Mon Sep 17 00:00:00 2001 From: Rajat Venkatesh <1638298+vrajat@users.noreply.github.com> Date: Wed, 23 Jul 2025 18:36:28 +0530 Subject: [PATCH 007/167] Allow Reset of ThreadResourceUsageAccountant in Tracing.java (#16360) --- .../broker/helix/BaseBrokerStarter.java | 5 +- .../helix/ControllerRequestClient.java | 21 +++++++ .../controller/helix/ControllerTest.java | 10 +++ .../CPUMemThreadLevelAccountingObjects.java | 1 + .../PerQueryCPUMemAccountantFactory.java | 15 ++++- .../ResourceUsageAccountantFactory.java | 11 +++- .../ResourceManagerAccountingTest.java | 61 +++++++++++-------- .../OOMProtectionEnabledIntegrationTest.java | 54 ++++++++++++++++ .../tests/WindowResourceAccountingTest.java | 1 + .../operator/MultiStageAccountingTest.java | 1 + ...MultistageResourceUsageAccountingTest.java | 1 + .../starter/helix/BaseServerStarter.java | 30 +++++---- .../ThreadResourceUsageAccountant.java | 7 +++ .../org/apache/pinot/spi/trace/Tracing.java | 51 +++++++++++++--- .../builder/ControllerRequestURLBuilder.java | 8 +++ 15 files changed, 223 insertions(+), 54 deletions(-) diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java index c694864055ee..3ae7fd846ce4 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java @@ -85,6 +85,7 @@ import org.apache.pinot.query.mailbox.MailboxService; import org.apache.pinot.query.service.dispatch.QueryDispatcher; import org.apache.pinot.segment.local.function.GroovyFunctionEvaluator; +import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant; import org.apache.pinot.spi.accounting.ThreadResourceUsageProvider; import org.apache.pinot.spi.cursors.ResponseStoreService; import org.apache.pinot.spi.env.PinotConfiguration; @@ -153,6 +154,7 @@ public abstract class BaseBrokerStarter implements ServiceStartable { protected AbstractResponseStore _responseStore; protected BrokerGrpcServer _brokerGrpcServer; protected FailureDetector _failureDetector; + protected ThreadResourceUsageAccountant _resourceUsageAccountant; @Override public void init(PinotConfiguration brokerConf) @@ -415,9 +417,10 @@ public void start() ThreadResourceUsageProvider.setThreadMemoryMeasurementEnabled( _brokerConf.getProperty(CommonConstants.Broker.CONFIG_OF_ENABLE_THREAD_ALLOCATED_BYTES_MEASUREMENT, CommonConstants.Broker.DEFAULT_THREAD_ALLOCATED_BYTES_MEASUREMENT)); - Tracing.ThreadAccountantOps.initializeThreadAccountant( + _resourceUsageAccountant = Tracing.ThreadAccountantOps.createThreadAccountant( _brokerConf.subset(CommonConstants.PINOT_QUERY_SCHEDULER_PREFIX), _instanceId, org.apache.pinot.spi.config.instance.InstanceType.BROKER); + Preconditions.checkNotNull(_resourceUsageAccountant); Tracing.ThreadAccountantOps.startThreadAccountant(); String controllerUrl = _brokerConf.getProperty(Broker.CONTROLLER_URL); diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/ControllerRequestClient.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/ControllerRequestClient.java index 8acf41df876b..c9efb0a41834 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/ControllerRequestClient.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/ControllerRequestClient.java @@ -473,4 +473,25 @@ protected static String getServerTenantRequestPayload(String tenantName, int num return new Tenant(TenantRole.SERVER, tenantName, numOfflineServers + numRealtimeServers, numOfflineServers, numRealtimeServers).toJsonString(); } + + public void updateClusterConfig(Map newConfigs) + throws IOException { + try { + HttpClient.wrapAndThrowHttpException(_httpClient.sendJsonPostRequest( + new URI(_controllerRequestURLBuilder.forClusterConfigUpdate()), + JsonUtils.objectToString(newConfigs), _headers)); + } catch (HttpErrorStatusException | URISyntaxException e) { + throw new IOException(e); + } + } + + public void deleteClusterConfig(String config) + throws IOException { + try { + HttpClient.wrapAndThrowHttpException(_httpClient.sendDeleteRequest( + new URI(_controllerRequestURLBuilder.forClusterConfigDelete(config)), _headers)); + } catch (HttpErrorStatusException | URISyntaxException e) { + throw new IOException(e); + } + } } diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/ControllerTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/ControllerTest.java index e3225dea0952..027b03e37b0a 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/ControllerTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/ControllerTest.java @@ -937,6 +937,16 @@ public void runPeriodicTask(String taskName, String tableName, TableType tableTy sendGetRequest(getControllerRequestURLBuilder().forPeriodTaskRun(taskName, tableName, tableType)); } + public void updateClusterConfig(Map clusterConfig) + throws IOException { + getControllerRequestClient().updateClusterConfig(clusterConfig); + } + + public void deleteClusterConfig(String clusterConfig) + throws IOException { + getControllerRequestClient().deleteClusterConfig(clusterConfig); + } + /** * Trigger a task on a table and wait for completion */ diff --git a/pinot-core/src/main/java/org/apache/pinot/core/accounting/CPUMemThreadLevelAccountingObjects.java b/pinot-core/src/main/java/org/apache/pinot/core/accounting/CPUMemThreadLevelAccountingObjects.java index 3744d1729b5a..b6a0bc661a9a 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/accounting/CPUMemThreadLevelAccountingObjects.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/accounting/CPUMemThreadLevelAccountingObjects.java @@ -76,6 +76,7 @@ public void setToIdle() { _currentThreadCPUTimeSampleMS = 0; // clear memory usage _currentThreadMemoryAllocationSampleBytes = 0; + _errorStatus.set(null); } /** diff --git a/pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java index 71944bd6108b..7402b5138f1c 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java @@ -90,7 +90,7 @@ public static class PerQueryCPUMemResourceUsageAccountant implements ThreadResou */ private static final String ACCOUNTANT_TASK_NAME = "CPUMemThreadAccountant"; private static final int ACCOUNTANT_PRIORITY = 4; - private static final ExecutorService EXECUTOR_SERVICE = Executors.newFixedThreadPool(1, r -> { + private final ExecutorService _executorService = Executors.newFixedThreadPool(1, r -> { Thread thread = new Thread(r); thread.setPriority(ACCOUNTANT_PRIORITY); thread.setDaemon(true); @@ -213,6 +213,10 @@ protected WatcherTask createWatcherTask() { return new WatcherTask(); } + public QueryMonitorConfig getQueryMonitorConfig() { + return _watcherTask.getQueryMonitorConfig(); + } + @Override public Collection getThreadResources() { return _threadEntriesMap.values(); @@ -431,7 +435,12 @@ public WatcherTask getWatcherTask() { @Override public void startWatcherTask() { - EXECUTOR_SERVICE.submit(_watcherTask); + _executorService.submit(_watcherTask); + } + + @Override + public void stopWatcherTask() { + _executorService.shutdownNow(); } @Override @@ -756,7 +765,7 @@ private void logQueryMonitorConfig() { @Override public void run() { - while (true) { + while (!Thread.currentThread().isInterrupted()) { try { runOnce(); } finally { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/accounting/ResourceUsageAccountantFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/accounting/ResourceUsageAccountantFactory.java index 4c4bcc8d0281..3b3d509bf5d7 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/accounting/ResourceUsageAccountantFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/accounting/ResourceUsageAccountantFactory.java @@ -54,7 +54,7 @@ public static class ResourceUsageAccountant implements ThreadResourceUsageAccoun private static final String ACCOUNTANT_TASK_NAME = "ResourceUsageAccountant"; private static final int ACCOUNTANT_PRIORITY = 4; - private static final ExecutorService EXECUTOR_SERVICE = Executors.newFixedThreadPool(1, r -> { + private final ExecutorService _executorService = Executors.newFixedThreadPool(1, r -> { Thread thread = new Thread(r); thread.setPriority(ACCOUNTANT_PRIORITY); thread.setDaemon(true); @@ -286,7 +286,12 @@ public void clear() { @Override public void startWatcherTask() { - EXECUTOR_SERVICE.submit(_watcherTask); + _executorService.submit(_watcherTask); + } + + @Override + public void stopWatcherTask() { + _executorService.shutdownNow(); } @Override @@ -315,7 +320,7 @@ class WatcherTask implements Runnable { @Override public void run() { LOGGER.debug("Running timed task for {}", this.getClass().getName()); - while (true) { + while (!Thread.currentThread().isInterrupted()) { try { // Preaggregation. runPreAggregation(); diff --git a/pinot-core/src/test/java/org/apache/pinot/core/accounting/ResourceManagerAccountingTest.java b/pinot-core/src/test/java/org/apache/pinot/core/accounting/ResourceManagerAccountingTest.java index 5e473e402bc7..4986e7684d35 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/accounting/ResourceManagerAccountingTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/accounting/ResourceManagerAccountingTest.java @@ -59,6 +59,7 @@ import org.apache.pinot.segment.spi.index.creator.JsonIndexCreator; import org.apache.pinot.segment.spi.memory.PinotDataBuffer; import org.apache.pinot.spi.accounting.ThreadExecutionContext; +import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant; import org.apache.pinot.spi.accounting.ThreadResourceUsageProvider; import org.apache.pinot.spi.config.instance.InstanceType; import org.apache.pinot.spi.config.table.JsonIndexConfig; @@ -99,14 +100,15 @@ public void testCPUtimeProvider() "org.apache.pinot.core.accounting.PerQueryCPUMemAccountantFactory"); configs.put(CommonConstants.Accounting.CONFIG_OF_ENABLE_THREAD_MEMORY_SAMPLING, false); configs.put(CommonConstants.Accounting.CONFIG_OF_ENABLE_THREAD_CPU_SAMPLING, true); - ResourceManager rm = getResourceManager(20, 40, 1, 1, configs); - Future[] futures = new Future[2000]; - AtomicInteger atomicInteger = new AtomicInteger(); PinotConfiguration pinotCfg = new PinotConfiguration(configs); - Tracing.ThreadAccountantOps.initializeThreadAccountant(pinotCfg, "testCPUtimeProvider", - InstanceType.SERVER); + ThreadResourceUsageAccountant accountant = Tracing.ThreadAccountantOps.createThreadAccountant(pinotCfg, + "testCPUtimeProvider", InstanceType.SERVER); Tracing.ThreadAccountantOps.startThreadAccountant(); + ResourceManager rm = getResourceManager(20, 40, 1, 1, configs, accountant); + Future[] futures = new Future[2000]; + AtomicInteger atomicInteger = new AtomicInteger(); + for (int k = 0; k < 30; k++) { int finalK = k; rm.getQueryRunners().submit(() -> { @@ -164,12 +166,13 @@ public void testThreadMemoryAccounting() "org.apache.pinot.core.accounting.PerQueryCPUMemAccountantFactory"); configs.put(CommonConstants.Accounting.CONFIG_OF_ENABLE_THREAD_MEMORY_SAMPLING, true); configs.put(CommonConstants.Accounting.CONFIG_OF_ENABLE_THREAD_CPU_SAMPLING, false); - ResourceManager rm = getResourceManager(20, 40, 1, 1, configs); PinotConfiguration pinotCfg = new PinotConfiguration(configs); - Tracing.ThreadAccountantOps.initializeThreadAccountant(pinotCfg, "testCPUtimeProvider", - InstanceType.SERVER); + ThreadResourceUsageAccountant accountant = Tracing.ThreadAccountantOps.createThreadAccountant(pinotCfg, + "testCPUtimeProvider", InstanceType.SERVER); Tracing.ThreadAccountantOps.startThreadAccountant(); + ResourceManager rm = getResourceManager(20, 40, 1, 1, configs, accountant); + for (int k = 0; k < 30; k++) { int finalK = k; rm.getQueryRunners().submit(() -> { @@ -241,13 +244,13 @@ public void testWorkloadLevelThreadMemoryAccounting() String workloadName = CommonConstants.Accounting.DEFAULT_WORKLOAD_NAME; PinotConfiguration pinotCfg = new PinotConfiguration(configs); - Tracing.ThreadAccountantOps.initializeThreadAccountant(pinotCfg, "testWorkloadMemoryAccounting", - InstanceType.SERVER); + ThreadResourceUsageAccountant accountant = Tracing.ThreadAccountantOps.createThreadAccountant(pinotCfg, + "testWorkloadMemoryAccounting", InstanceType.SERVER); Tracing.ThreadAccountantOps.startThreadAccountant(); WorkloadBudgetManager workloadBudgetManager = Tracing.ThreadAccountantOps.getWorkloadBudgetManager(); workloadBudgetManager.addOrUpdateWorkload(workloadName, 88_000_000, 27_000_000); - ResourceManager rm = getResourceManager(20, 40, 1, 1, configs); + ResourceManager rm = getResourceManager(20, 40, 1, 1, configs, accountant); for (int k = 0; k < 30; k++) { int finalK = k; @@ -301,7 +304,8 @@ public void testWorkloadLevelThreadMemoryAccounting() @Test public void testWorkerThreadInterruption() throws Exception { - ResourceManager rm = getResourceManager(2, 5, 1, 3, Collections.emptyMap()); + ResourceManager rm = getResourceManager(2, 5, 1, 3, Collections.emptyMap(), + new Tracing.DefaultThreadResourceUsageAccountant()); AtomicReference[] futures = new AtomicReference[5]; for (int i = 0; i < 5; i++) { futures[i] = new AtomicReference<>(); @@ -379,10 +383,12 @@ public void testGetDataTableOOMSelect(String accountantFactoryClass) configs.put(CommonConstants.Accounting.CONFIG_OF_ENABLE_THREAD_CPU_SAMPLING, false); configs.put(CommonConstants.Accounting.CONFIG_OF_OOM_PROTECTION_KILLING_QUERY, true); PinotConfiguration config = getConfig(20, 2, configs); - ResourceManager rm = getResourceManager(20, 2, 1, 1, configs); // init accountant and start watcher task - Tracing.ThreadAccountantOps.initializeThreadAccountant(config, "testSelect", InstanceType.SERVER); + Tracing.unregisterThreadAccountant(); + ThreadResourceUsageAccountant accountant = Tracing.ThreadAccountantOps.createThreadAccountant(config, + "testSelect", InstanceType.SERVER); Tracing.ThreadAccountantOps.startThreadAccountant(); + ResourceManager rm = getResourceManager(20, 2, 1, 1, configs, accountant); CountDownLatch latch = new CountDownLatch(100); AtomicBoolean earlyTerminationOccurred = new AtomicBoolean(false); @@ -404,7 +410,7 @@ public void testGetDataTableOOMSelect(String accountantFactoryClass) } }); } - latch.await(); + latch.await(1, java.util.concurrent.TimeUnit.MINUTES); // assert that EarlyTerminationException was thrown in at least one runner thread Assert.assertTrue(earlyTerminationOccurred.get()); } @@ -450,11 +456,14 @@ public void testGetDataTableOOMGroupBy(String accountantFactoryClass) configs.put(CommonConstants.Accounting.CONFIG_OF_OOM_PROTECTION_KILLING_QUERY, true); PinotConfiguration config = getConfig(20, 2, configs); - ResourceManager rm = getResourceManager(20, 2, 1, 1, configs); // init accountant and start watcher task - Tracing.ThreadAccountantOps.initializeThreadAccountant(config, "testGroupBy", InstanceType.SERVER); + Tracing.unregisterThreadAccountant(); + ThreadResourceUsageAccountant accountant = Tracing.ThreadAccountantOps.createThreadAccountant(config, + "testGroupBy", InstanceType.SERVER); Tracing.ThreadAccountantOps.startThreadAccountant(); + ResourceManager rm = getResourceManager(20, 2, 1, 1, configs, accountant); + CountDownLatch latch = new CountDownLatch(100); AtomicBoolean earlyTerminationOccurred = new AtomicBoolean(false); @@ -475,7 +484,7 @@ public void testGetDataTableOOMGroupBy(String accountantFactoryClass) } }); } - latch.await(); + latch.await(1, java.util.concurrent.TimeUnit.MINUTES); // assert that EarlyTerminationException was thrown in at least one runner thread Assert.assertTrue(earlyTerminationOccurred.get()); } @@ -507,12 +516,14 @@ public void testJsonIndexExtractMapOOM(String accountantFactoryClass) configs.put(CommonConstants.Accounting.CONFIG_OF_MIN_MEMORY_FOOTPRINT_TO_KILL_RATIO, 0.00f); PinotConfiguration config = getConfig(2, 2, configs); - ResourceManager rm = getResourceManager(2, 2, 1, 1, configs); // init accountant and start watcher task - Tracing.ThreadAccountantOps.initializeThreadAccountant(config, "testJsonIndexExtractMapOOM", - InstanceType.SERVER); + Tracing.unregisterThreadAccountant(); + ThreadResourceUsageAccountant accountant = Tracing.ThreadAccountantOps.createThreadAccountant(config, + "testJsonIndexExtractMapOOM", InstanceType.SERVER); Tracing.ThreadAccountantOps.startThreadAccountant(); + ResourceManager rm = getResourceManager(2, 2, 1, 1, configs, accountant); + Supplier randomJsonValue = () -> { Random random = new Random(); int length = random.nextInt(1000); @@ -575,7 +586,7 @@ public void testJsonIndexExtractMapOOM(String accountantFactoryClass) } }); - latch.await(); + latch.await(1, java.util.concurrent.TimeUnit.MINUTES); Assert.assertTrue(mutableEarlyTerminationOccurred.get(), "Expected early termination reading the mutable index"); Assert.assertTrue(immutableEarlyTerminationOccurred.get(), @@ -602,7 +613,7 @@ public void testThreadMemory() "org.apache.pinot.core.accounting.PerQueryCPUMemAccountantFactory"); configs.put(CommonConstants.Accounting.CONFIG_OF_ENABLE_THREAD_MEMORY_SAMPLING, true); configs.put(CommonConstants.Accounting.CONFIG_OF_ENABLE_THREAD_CPU_SAMPLING, false); - ResourceManager rm = getResourceManager(20, 40, 1, 1, configs); + ResourceManager rm = getResourceManager(20, 40, 1, 1, configs, new Tracing.DefaultThreadResourceUsageAccountant()); Future[] futures = new Future[30]; for (int k = 0; k < 4; k++) { @@ -651,9 +662,9 @@ public void testThreadMemory() } private ResourceManager getResourceManager(int runners, int workers, final int softLimit, final int hardLimit, - Map map) { + Map map, ThreadResourceUsageAccountant accountant) { - return new ResourceManager(getConfig(runners, workers, map), new Tracing.DefaultThreadResourceUsageAccountant()) { + return new ResourceManager(getConfig(runners, workers, map), accountant) { @Override public QueryExecutorService getExecutorService(ServerQueryRequest query, SchedulerGroupAccountant accountant) { diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OOMProtectionEnabledIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OOMProtectionEnabledIntegrationTest.java index 68c3a8a42165..0ad70dac9be1 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OOMProtectionEnabledIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OOMProtectionEnabledIntegrationTest.java @@ -19,15 +19,23 @@ package org.apache.pinot.integration.tests; import java.io.File; +import java.io.IOException; import java.util.List; +import java.util.Map; +import org.apache.pinot.core.accounting.PerQueryCPUMemAccountantFactory; +import org.apache.pinot.core.accounting.QueryMonitorConfig; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.data.Schema; import org.apache.pinot.spi.env.PinotConfiguration; +import org.apache.pinot.spi.trace.Tracing; import org.apache.pinot.spi.utils.CommonConstants; import org.apache.pinot.util.TestUtils; import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; + public class OOMProtectionEnabledIntegrationTest extends BaseClusterIntegrationTestSet { private static final int NUM_OFFLINE_SEGMENTS = 8; @@ -62,6 +70,7 @@ public void setUp() startZk(); startController(); startBroker(); + Tracing.unregisterThreadAccountant(); startServer(); startKafka(); @@ -100,4 +109,49 @@ public void testHardcodedQueries(boolean useMultiStageEngine) setUseMultiStageQueryEngine(useMultiStageEngine); super.testHardcodedQueries(); } + + @Test + public void testChangeOomKillQueryEnabled() + throws IOException { + assertTrue(_serverStarters.get(0) + .getResourceUsageAccountant() instanceof PerQueryCPUMemAccountantFactory.PerQueryCPUMemResourceUsageAccountant); + PerQueryCPUMemAccountantFactory.PerQueryCPUMemResourceUsageAccountant accountant = + (PerQueryCPUMemAccountantFactory.PerQueryCPUMemResourceUsageAccountant) _serverStarters.get(0) + .getResourceUsageAccountant(); + + QueryMonitorConfig queryMonitorConfig = accountant.getQueryMonitorConfig(); + assertFalse(queryMonitorConfig.isOomKillQueryEnabled()); + + updateClusterConfig(Map.of("pinot.query.scheduler.accounting.oom.enable.killing.query", "true", + "pinot.query.scheduler.accounting.query.killed.metric.enabled", "true")); + + TestUtils.waitForCondition(aVoid -> { + QueryMonitorConfig updatedQueryMonitorConfig = accountant.getQueryMonitorConfig(); + return updatedQueryMonitorConfig.isOomKillQueryEnabled() + && updatedQueryMonitorConfig.isQueryKilledMetricEnabled(); + }, 1000L, "Waiting for OOM protection to be enabled"); + } + + @Test + public void testChangeThresholds() + throws IOException { + assertTrue(_serverStarters.get(0) + .getResourceUsageAccountant() instanceof PerQueryCPUMemAccountantFactory.PerQueryCPUMemResourceUsageAccountant); + PerQueryCPUMemAccountantFactory.PerQueryCPUMemResourceUsageAccountant accountant = + (PerQueryCPUMemAccountantFactory.PerQueryCPUMemResourceUsageAccountant) _serverStarters.get(0) + .getResourceUsageAccountant(); + + + QueryMonitorConfig queryMonitorConfig = accountant.getQueryMonitorConfig(); + updateClusterConfig(Map.of("pinot.query.scheduler.accounting.oom.alarming.usage.ratio", "0.7f", + "pinot.query.scheduler.accounting.oom.critical.heap.usage.ratio", "0.75f", + "pinot.query.scheduler.accounting.oom.panic.heap.usage.ratio", "0.8f")); + + TestUtils.waitForCondition(aVoid -> { + QueryMonitorConfig updatedQueryMonitorConfig = accountant.getQueryMonitorConfig(); + return updatedQueryMonitorConfig.getAlarmingLevel() != queryMonitorConfig.getAlarmingLevel() + && updatedQueryMonitorConfig.getCriticalLevel() != queryMonitorConfig.getCriticalLevel() + && updatedQueryMonitorConfig.getPanicLevel() != queryMonitorConfig.getPanicLevel(); + }, 1000L, "Waiting for OOM protection to be enabled"); + } } diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/WindowResourceAccountingTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/WindowResourceAccountingTest.java index d511787f8eb6..b6158211a021 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/WindowResourceAccountingTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/WindowResourceAccountingTest.java @@ -83,6 +83,7 @@ public void setUp() startZk(); startController(); startBroker(); + Tracing.unregisterThreadAccountant(); startServer(); if (_controllerRequestURLBuilder == null) { diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MultiStageAccountingTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MultiStageAccountingTest.java index 37a21e314fd7..de34a977d05c 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MultiStageAccountingTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MultiStageAccountingTest.java @@ -91,6 +91,7 @@ public static void setUpClass() { configs.put(CommonConstants.Accounting.CONFIG_OF_ENABLE_THREAD_CPU_SAMPLING, false); configs.put(CommonConstants.Accounting.CONFIG_OF_OOM_PROTECTION_KILLING_QUERY, true); // init accountant and start watcher task + Tracing.unregisterThreadAccountant(); Tracing.ThreadAccountantOps.initializeThreadAccountant(new PinotConfiguration(configs), "testGroupBy", InstanceType.SERVER); Tracing.ThreadAccountantOps.startThreadAccountant(); diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MultistageResourceUsageAccountingTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MultistageResourceUsageAccountingTest.java index 9f19b55171af..534f1f413d54 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MultistageResourceUsageAccountingTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MultistageResourceUsageAccountingTest.java @@ -91,6 +91,7 @@ public static void setUpClass() { configs.put(CommonConstants.Accounting.CONFIG_OF_WORKLOAD_ENABLE_COST_COLLECTION, true); // init accountant and start watcher task PinotConfiguration pinotCfg = new PinotConfiguration(configs); + Tracing.unregisterThreadAccountant(); Tracing.ThreadAccountantOps.initializeThreadAccountant(pinotCfg, "testGroupBy", InstanceType.SERVER); Tracing.ThreadAccountantOps.startThreadAccountant(); diff --git a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/BaseServerStarter.java b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/BaseServerStarter.java index 27e6ecc8524c..c3f6f0e611ad 100644 --- a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/BaseServerStarter.java +++ b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/BaseServerStarter.java @@ -172,6 +172,7 @@ public abstract class BaseServerStarter implements ServiceStartable { protected DefaultClusterConfigChangeHandler _clusterConfigChangeHandler; protected volatile boolean _isServerReadyToServeQueries = false; private ScheduledExecutorService _helixMessageCountScheduler; + protected ThreadResourceUsageAccountant _resourceUsageAccountant; @Override public void init(PinotConfiguration serverConf) @@ -243,14 +244,6 @@ public void init(PinotConfiguration serverConf) // Initialize the data buffer factory PinotDataBuffer.loadDefaultFactory(serverConf); - // Enable/disable thread CPU time measurement through instance config. - ThreadResourceUsageProvider.setThreadCpuTimeMeasurementEnabled( - _serverConf.getProperty(Server.CONFIG_OF_ENABLE_THREAD_CPU_TIME_MEASUREMENT, - Server.DEFAULT_ENABLE_THREAD_CPU_TIME_MEASUREMENT)); - // Enable/disable thread memory allocation tracking through instance config - ThreadResourceUsageProvider.setThreadMemoryMeasurementEnabled( - _serverConf.getProperty(Server.CONFIG_OF_ENABLE_THREAD_ALLOCATED_BYTES_MEASUREMENT, - Server.DEFAULT_THREAD_ALLOCATED_BYTES_MEASUREMENT)); // Set data table version send to broker. int dataTableVersion = _serverConf.getProperty(Server.CONFIG_OF_CURRENT_DATA_TABLE_VERSION, DataTableBuilderFactory.DEFAULT_VERSION); @@ -670,16 +663,24 @@ public void start() segmentDownloadThrottler, segmentMultiColTextIndexPreprocessThrottler); } + // Enable/disable thread CPU time measurement through instance config. + ThreadResourceUsageProvider.setThreadCpuTimeMeasurementEnabled( + _serverConf.getProperty(Server.CONFIG_OF_ENABLE_THREAD_CPU_TIME_MEASUREMENT, + Server.DEFAULT_ENABLE_THREAD_CPU_TIME_MEASUREMENT)); + // Enable/disable thread memory allocation tracking through instance config + ThreadResourceUsageProvider.setThreadMemoryMeasurementEnabled( + _serverConf.getProperty(Server.CONFIG_OF_ENABLE_THREAD_ALLOCATED_BYTES_MEASUREMENT, + Server.DEFAULT_THREAD_ALLOCATED_BYTES_MEASUREMENT)); // Initialize the thread accountant for query killing - Tracing.ThreadAccountantOps.initializeThreadAccountant( + _resourceUsageAccountant = Tracing.ThreadAccountantOps.createThreadAccountant( _serverConf.subset(CommonConstants.PINOT_QUERY_SCHEDULER_PREFIX), _instanceId, org.apache.pinot.spi.config.instance.InstanceType.SERVER); - ThreadResourceUsageAccountant threadAccountant = Tracing.getThreadAccountant(); + Preconditions.checkNotNull(_resourceUsageAccountant); SendStatsPredicate sendStatsPredicate = SendStatsPredicate.create(_serverConf, _helixManager); ServerConf serverConf = new ServerConf(_serverConf); _serverInstance = new ServerInstance(serverConf, _helixManager, _accessControlFactory, _segmentOperationsThrottler, - sendStatsPredicate, threadAccountant); + sendStatsPredicate, _resourceUsageAccountant); ServerMetrics serverMetrics = _serverInstance.getServerMetrics(); InstanceDataManager instanceDataManager = _serverInstance.getInstanceDataManager(); @@ -784,7 +785,8 @@ public void start() // Start the thread accountant Tracing.ThreadAccountantOps.startThreadAccountant(); - PinotClusterConfigChangeListener threadAccountantListener = threadAccountant.getClusterConfigChangeListener(); + PinotClusterConfigChangeListener threadAccountantListener = + _resourceUsageAccountant.getClusterConfigChangeListener(); if (threadAccountantListener != null) { _clusterConfigChangeHandler.registerClusterConfigChangeListener(threadAccountantListener); } @@ -1149,4 +1151,8 @@ private void refreshMessageCount() { LOGGER.warn("Failed to refresh Helix message count", e); } } + + public ThreadResourceUsageAccountant getResourceUsageAccountant() { + return _resourceUsageAccountant; + } } diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceUsageAccountant.java b/pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceUsageAccountant.java index abf823fa6f29..4fdb4f1e1b60 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceUsageAccountant.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceUsageAccountant.java @@ -132,6 +132,13 @@ void updateQueryUsageConcurrently(String identifier, long cpuTimeNs, long alloca */ void startWatcherTask(); + /** + * Stop the periodic watcher task. + */ + default void stopWatcherTask() { + // Default implementation does nothing. Subclasses can override to stop the watcher task. + } + @Nullable default PinotClusterConfigChangeListener getClusterConfigChangeListener() { return null; diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/trace/Tracing.java b/pinot-spi/src/main/java/org/apache/pinot/spi/trace/Tracing.java index 4cb83ab79c09..afc9a283c1b8 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/trace/Tracing.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/trace/Tracing.java @@ -67,7 +67,7 @@ private Tracing() { private static final class Holder { static final Tracer TRACER = TRACER_REGISTRATION.get() == null ? createDefaultTracer() : TRACER_REGISTRATION.get(); - static final ThreadResourceUsageAccountant ACCOUNTANT = + static ThreadResourceUsageAccountant _accountant = ACCOUNTANT_REGISTRATION.get() == null ? createDefaultThreadAccountant() : ACCOUNTANT_REGISTRATION.get(); } @@ -88,7 +88,12 @@ public static boolean register(Tracer tracer) { * @return true if the registration was successful. */ public static boolean register(ThreadResourceUsageAccountant threadResourceUsageAccountant) { - return ACCOUNTANT_REGISTRATION.compareAndSet(null, threadResourceUsageAccountant); + if (ACCOUNTANT_REGISTRATION.compareAndSet(null, threadResourceUsageAccountant)) { + Holder._accountant = threadResourceUsageAccountant; + LOGGER.info("Registered thread accountant: {}", threadResourceUsageAccountant.getClass().getName()); + return true; + } + return false; } /** @@ -109,7 +114,7 @@ public static Tracer getTracer() { * @return the registered threadAccountant. */ public static ThreadResourceUsageAccountant getThreadAccountant() { - return Holder.ACCOUNTANT; + return Holder._accountant; } /** @@ -138,7 +143,23 @@ private static Tracer createDefaultTracer() { */ private static DefaultThreadResourceUsageAccountant createDefaultThreadAccountant() { LOGGER.info("Using default thread accountant"); - return new DefaultThreadResourceUsageAccountant(); + DefaultThreadResourceUsageAccountant accountant = new DefaultThreadResourceUsageAccountant(); + Holder._accountant = accountant; + ACCOUNTANT_REGISTRATION.set(accountant); + return accountant; + } + + /** + * Unregisters the thread accountant. This is only used in tests when a custom thread accountant is required. + * This will reset the thread accountant to null, so that the next call to initializeThreadAccountant or + * createThreadAccountant will register the new thread accountant. + */ + public static void unregisterThreadAccountant() { + if (Holder._accountant != null) { + Holder._accountant.stopWatcherTask(); + } + Holder._accountant = null; + ACCOUNTANT_REGISTRATION.set(null); } /** @@ -323,25 +344,35 @@ public static void clear() { public static void initializeThreadAccountant(PinotConfiguration config, String instanceId, InstanceType instanceType) { - String factoryName = config.getProperty(CommonConstants.Accounting.CONFIG_OF_FACTORY_NAME); + createThreadAccountant(config, instanceId, instanceType); + } + + public static ThreadResourceUsageAccountant createThreadAccountant(PinotConfiguration config, String instanceId, + InstanceType instanceType) { _workloadBudgetManager = new WorkloadBudgetManager(config); - if (factoryName == null) { - LOGGER.warn("No thread accountant factory provided, using default implementation"); - } else { + String factoryName = config.getProperty(CommonConstants.Accounting.CONFIG_OF_FACTORY_NAME); + ThreadResourceUsageAccountant accountant = null; + if (factoryName != null) { LOGGER.info("Config-specified accountant factory name {}", factoryName); try { ThreadAccountantFactory threadAccountantFactory = (ThreadAccountantFactory) Class.forName(factoryName).getDeclaredConstructor().newInstance(); - boolean registered = Tracing.register(threadAccountantFactory.init(config, instanceId, instanceType)); LOGGER.info("Using accountant provided by {}", factoryName); + accountant = threadAccountantFactory.init(config, instanceId, instanceType); + boolean registered = register(accountant); if (!registered) { - LOGGER.warn("ThreadAccountant {} register unsuccessful, as it is already registered.", factoryName); + LOGGER.warn("ThreadAccountant register unsuccessful, as it is already registered."); } } catch (Exception exception) { LOGGER.warn("Using default implementation of thread accountant, " + "due to invalid thread accountant factory {} provided, exception:", factoryName, exception); } } + // If no factory is specified or the factory creation failed, use the default implementation + if (accountant == null) { + accountant = createDefaultThreadAccountant(); + } + return accountant; } public static void startThreadAccountant() { diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/builder/ControllerRequestURLBuilder.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/builder/ControllerRequestURLBuilder.java index cb5de2956ace..e90e4272398e 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/builder/ControllerRequestURLBuilder.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/builder/ControllerRequestURLBuilder.java @@ -694,4 +694,12 @@ public String forLogicalTableDelete(String logicalTableName) { public String forTableTimeBoundary(String tableName) { return StringUtil.join("/", _baseUrl, "tables", tableName, "timeBoundary"); } + + public String forClusterConfigUpdate() { + return StringUtil.join("/", _baseUrl, "cluster", "configs"); + } + + public String forClusterConfigDelete(String config) { + return StringUtil.join("/", _baseUrl, "cluster", "configs", config); + } } From c1c891602766a5a1de2765bd4fa5b37788a6f846 Mon Sep 17 00:00:00 2001 From: 9aman <35227405+9aman@users.noreply.github.com> Date: Thu, 24 Jul 2025 01:24:51 +0530 Subject: [PATCH 008/167] Cleanup logs (#16410) --- .../pinot/core/accounting/PerQueryCPUMemAccountantFactory.java | 2 -- 1 file changed, 2 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java index 7402b5138f1c..e2f23664e392 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java @@ -914,8 +914,6 @@ void killAllQueries() { /** * Kill the query with the highest cost (memory footprint/cpu time/...) - * Will trigger gc when killing a consecutive number of queries - * use XX:+ExplicitGCInvokesConcurrent to avoid a full gc when system.gc is triggered */ private void killMostExpensiveQuery() { if (!_isThreadMemorySamplingEnabled) { From c7dd57c65623229408eaa7c644020b81bf1b23ba Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Wed, 23 Jul 2025 14:29:26 -0700 Subject: [PATCH 009/167] Enhance LookupTable implementations and add comprehensive test coverage for join/window operators (#16408) - Add size() method to all LookupTable implementations: * IntLookupTable, LongLookupTable, FloatLookupTable, DoubleLookupTable, ObjectLookupTable, LookupTable base class - Update join and window operators to utilize LookupTable.size() for performance tracking: * HashJoinOperator: Use size() for thread accounting and right table size tracking * AsofJoinOperator: Add size-based performance monitoring * NonEquiJoinOperator: Integrate size() for resource accounting * WindowAggregateOperator: Add lookup table size tracking - Add comprehensive LookupTableTest suite (484 lines, 47 test cases): * Test coverage for all LookupTable implementations * Empty table behavior, unique/duplicate key scenarios, null key handling * Mixed scenarios, finish() behavior, type-specific operations * SQL-compliant semantics and edge case validation - All tests passing with enhanced join operator infrastructure --- .../runtime/operator/AsofJoinOperator.java | 5 + .../runtime/operator/HashJoinOperator.java | 14 +- .../runtime/operator/NonEquiJoinOperator.java | 4 + .../operator/WindowAggregateOperator.java | 4 + .../operator/join/DoubleLookupTable.java | 5 + .../operator/join/FloatLookupTable.java | 5 + .../runtime/operator/join/IntLookupTable.java | 5 + .../operator/join/LongLookupTable.java | 5 + .../runtime/operator/join/LookupTable.java | 5 + .../operator/join/ObjectLookupTable.java | 5 + .../operator/join/LookupTableTest.java | 484 ++++++++++++++++++ 11 files changed, 539 insertions(+), 2 deletions(-) create mode 100644 pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/join/LookupTableTest.java diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AsofJoinOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AsofJoinOperator.java index 3dd0a7781e0f..73d229a23042 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AsofJoinOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/AsofJoinOperator.java @@ -32,6 +32,7 @@ import org.apache.pinot.query.planner.plannode.JoinNode; import org.apache.pinot.query.runtime.blocks.MseBlock; import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; +import org.apache.pinot.spi.trace.Tracing; public class AsofJoinOperator extends BaseJoinOperator { @@ -102,6 +103,7 @@ protected List buildJoinedRows(MseBlock.Data leftBlock) { if (matchKey == null) { // Rows with null match keys cannot be matched with any right rows if (needUnmatchedLeftRows()) { + Tracing.ThreadAccountantOps.sampleAndCheckInterruptionPeriodically(rows.size()); rows.add(joinRow(leftRow, null)); } continue; @@ -110,15 +112,18 @@ protected List buildJoinedRows(MseBlock.Data leftBlock) { NavigableMap, Object[]> rightRows = _rightTable.get(hashKey); if (rightRows == null) { if (needUnmatchedLeftRows()) { + Tracing.ThreadAccountantOps.sampleAndCheckInterruptionPeriodically(rows.size()); rows.add(joinRow(leftRow, null)); } } else { Object[] rightRow = closestMatch(matchKey, rightRows); if (rightRow == null) { if (needUnmatchedLeftRows()) { + Tracing.ThreadAccountantOps.sampleAndCheckInterruptionPeriodically(rows.size()); rows.add(joinRow(leftRow, null)); } } else { + Tracing.ThreadAccountantOps.sampleAndCheckInterruptionPeriodically(rows.size()); rows.add(joinRow(leftRow, rightRow)); } } diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java index 525c426ed60b..0a33fcd12ac6 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java @@ -39,6 +39,7 @@ import org.apache.pinot.query.runtime.operator.join.LookupTable; import org.apache.pinot.query.runtime.operator.join.ObjectLookupTable; import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; +import org.apache.pinot.spi.trace.Tracing; /** @@ -114,6 +115,7 @@ protected void addRowsToRightTable(List rows) { } continue; } + Tracing.ThreadAccountantOps.sampleAndCheckInterruptionPeriodically(_rightTable.size()); _rightTable.addRow(key, row); } } @@ -138,8 +140,6 @@ private boolean isNullKey(Object key) { return false; } - - @Override protected void finishBuildingRightTable() { assert _rightTable != null : "Right table should not be null when finishing building"; @@ -203,6 +203,7 @@ private List buildJoinedDataBlockUniqueKeys(MseBlock.Data leftBlock) { break; } // defer copying of the content until row matches + Tracing.ThreadAccountantOps.sampleAndCheckInterruptionPeriodically(rows.size()); rows.add(resultRowView.toArray()); if (_matchedRightRows != null) { _matchedRightRows.put(key, BIT_SET_PLACEHOLDER); @@ -222,6 +223,7 @@ private List buildJoinedDataBlockDuplicateKeys(MseBlock.Data leftBlock List rows = new ArrayList<>(leftRows.size()); for (Object[] leftRow : leftRows) { + Tracing.ThreadAccountantOps.sampleAndCheckInterruptionPeriodically(rows.size()); Object key = _leftKeySelector.getKey(leftRow); // Skip rows with null join keys - they should not participate in equi-joins per SQL standard if (handleNullKey(key, leftRow, rows)) { @@ -241,6 +243,7 @@ private List buildJoinedDataBlockDuplicateKeys(MseBlock.Data leftBlock maxRowsLimitReached = true; break; } + Tracing.ThreadAccountantOps.sampleAndCheckInterruptionPeriodically(rows.size()); rows.add(resultRowView.toArray()); hasMatchForLeftRow = true; if (_matchedRightRows != null) { @@ -265,6 +268,7 @@ private void handleUnmatchedLeftRow(Object[] leftRow, List rows) { if (isMaxRowsLimitReached(rows.size())) { return; } + Tracing.ThreadAccountantOps.sampleAndCheckInterruptionPeriodically(rows.size()); rows.add(joinRow(leftRow, null)); } } @@ -277,6 +281,7 @@ private List buildJoinedDataBlockSemi(MseBlock.Data leftBlock) { for (Object[] leftRow : leftRows) { Object key = _leftKeySelector.getKey(leftRow); if (_rightTable.containsKey(key)) { + Tracing.ThreadAccountantOps.sampleAndCheckInterruptionPeriodically(rows.size()); rows.add(leftRow); } } @@ -292,6 +297,7 @@ private List buildJoinedDataBlockAnti(MseBlock.Data leftBlock) { for (Object[] leftRow : leftRows) { Object key = _leftKeySelector.getKey(leftRow); if (!_rightTable.containsKey(key)) { + Tracing.ThreadAccountantOps.sampleAndCheckInterruptionPeriodically(rows.size()); rows.add(leftRow); } } @@ -308,6 +314,7 @@ protected List buildNonMatchRightRows() { for (Map.Entry entry : _rightTable.entrySet()) { Object[] rightRow = (Object[]) entry.getValue(); if (!_matchedRightRows.containsKey(entry.getKey())) { + Tracing.ThreadAccountantOps.sampleAndCheckInterruptionPeriodically(rows.size()); rows.add(joinRow(null, rightRow)); } } @@ -317,12 +324,14 @@ protected List buildNonMatchRightRows() { BitSet matchedIndices = _matchedRightRows.get(entry.getKey()); if (matchedIndices == null) { for (Object[] rightRow : rightRows) { + Tracing.ThreadAccountantOps.sampleAndCheckInterruptionPeriodically(rows.size()); rows.add(joinRow(null, rightRow)); } } else { int numRightRows = rightRows.size(); int unmatchedIndex = 0; while ((unmatchedIndex = matchedIndices.nextClearBit(unmatchedIndex)) < numRightRows) { + Tracing.ThreadAccountantOps.sampleAndCheckInterruptionPeriodically(rows.size()); rows.add(joinRow(null, rightRows.get(unmatchedIndex++))); } } @@ -331,6 +340,7 @@ protected List buildNonMatchRightRows() { // Add unmatched null key rows from right side for RIGHT and FULL JOIN if (_nullKeyRightRows != null) { for (Object[] nullKeyRow : _nullKeyRightRows) { + Tracing.ThreadAccountantOps.sampleAndCheckInterruptionPeriodically(rows.size()); rows.add(joinRow(null, nullKeyRow)); } } diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/NonEquiJoinOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/NonEquiJoinOperator.java index 638d4d98806a..6a5a15e97ac7 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/NonEquiJoinOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/NonEquiJoinOperator.java @@ -28,6 +28,7 @@ import org.apache.pinot.query.planner.plannode.JoinNode; import org.apache.pinot.query.runtime.blocks.MseBlock; import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; +import org.apache.pinot.spi.trace.Tracing; /** @@ -90,6 +91,7 @@ protected List buildJoinedRows(MseBlock.Data leftBlock) { maxRowsLimitReached = true; break; } + Tracing.ThreadAccountantOps.sampleAndCheckInterruptionPeriodically(rows.size()); rows.add(joinRowView.toArray()); hasMatchForLeftRow = true; if (_matchedRightRows != null) { @@ -104,6 +106,7 @@ protected List buildJoinedRows(MseBlock.Data leftBlock) { if (isMaxRowsLimitReached(rows.size())) { break; } + Tracing.ThreadAccountantOps.sampleAndCheckInterruptionPeriodically(rows.size()); rows.add(joinRow(leftRow, null)); } } @@ -121,6 +124,7 @@ protected List buildNonMatchRightRows() { List rows = new ArrayList<>(numRightRows - numMatchedRightRows); int unmatchedIndex = 0; while ((unmatchedIndex = _matchedRightRows.nextClearBit(unmatchedIndex)) < numRightRows) { + Tracing.ThreadAccountantOps.sampleAndCheckInterruptionPeriodically(rows.size()); rows.add(joinRow(null, _rightTable.get(unmatchedIndex++))); } return rows; diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperator.java index f50fe2a5db74..51dd28122614 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/WindowAggregateOperator.java @@ -43,6 +43,7 @@ import org.apache.pinot.query.runtime.operator.window.WindowFunctionFactory; import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; import org.apache.pinot.spi.exception.QueryErrorCode; +import org.apache.pinot.spi.trace.Tracing; import org.apache.pinot.spi.utils.CommonConstants.MultiStageQueryRunner.WindowOverFlowMode; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -230,6 +231,7 @@ private MseBlock computeBlocks() { for (Object[] row : container) { // TODO: Revisit null direction handling for all query types Key key = AggregationUtils.extractRowKey(row, _keys); + Tracing.ThreadAccountantOps.sampleAndCheckInterruptionPeriodically(_numRows); partitionRows.computeIfAbsent(key, k -> new ArrayList<>()).add(row); } _numRows += containerSize; @@ -253,6 +255,7 @@ private MseBlock computeBlocks() { for (WindowFunction windowFunction : _windowFunctions) { List processRows = windowFunction.processRows(rowList); assert processRows.size() == rowList.size(); + Tracing.ThreadAccountantOps.sampleAndCheckInterruptionPeriodically(windowFunctionResults.size()); windowFunctionResults.add(processRows); } @@ -265,6 +268,7 @@ private MseBlock computeBlocks() { } // Convert the results from WindowFunction to the desired type TypeUtils.convertRow(row, resultStoredTypes); + Tracing.ThreadAccountantOps.sampleAndCheckInterruptionPeriodically(rows.size()); rows.add(row); } } diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/DoubleLookupTable.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/DoubleLookupTable.java index 3ca59141e218..0e4d76fc4760 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/DoubleLookupTable.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/DoubleLookupTable.java @@ -62,4 +62,9 @@ public Object lookupNotNullKey(Object key) { public Set> notNullKeyEntrySet() { return (Set) _lookupTable.double2ObjectEntrySet(); } + + @Override + public int size() { + return _lookupTable.size(); + } } diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/FloatLookupTable.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/FloatLookupTable.java index ae7a2c106c43..79a815d5a2be 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/FloatLookupTable.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/FloatLookupTable.java @@ -62,4 +62,9 @@ public Object lookupNotNullKey(Object key) { public Set> notNullKeyEntrySet() { return (Set) _lookupTable.float2ObjectEntrySet(); } + + @Override + public int size() { + return _lookupTable.size(); + } } diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/IntLookupTable.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/IntLookupTable.java index fb6ef693a005..8656b246a569 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/IntLookupTable.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/IntLookupTable.java @@ -57,4 +57,9 @@ protected Object lookupNotNullKey(Object key) { protected Set> notNullKeyEntrySet() { return (Set) _lookupTable.int2ObjectEntrySet(); } + + @Override + public int size() { + return _lookupTable.size(); + } } diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/LongLookupTable.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/LongLookupTable.java index 44fc2ed61def..9c486ebdbc7f 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/LongLookupTable.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/LongLookupTable.java @@ -62,4 +62,9 @@ public Object lookupNotNullKey(Object key) { public Set> notNullKeyEntrySet() { return (Set) _lookupTable.long2ObjectEntrySet(); } + + @Override + public int size() { + return _lookupTable.size(); + } } diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/LookupTable.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/LookupTable.java index 86f811eca56f..bfaacd301831 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/LookupTable.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/LookupTable.java @@ -104,4 +104,9 @@ public boolean isKeysUnique() { */ @SuppressWarnings("rawtypes") public abstract Set> entrySet(); + + /** + * Returns the number of entries in the lookup table. + */ + public abstract int size(); } diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/ObjectLookupTable.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/ObjectLookupTable.java index d94ac1cb84ac..46bf29ec89b2 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/ObjectLookupTable.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/ObjectLookupTable.java @@ -71,4 +71,9 @@ public Object lookup(@Nullable Object key) { public Set> entrySet() { return _lookupTable.entrySet(); } + + @Override + public int size() { + return _lookupTable.size(); + } } diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/join/LookupTableTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/join/LookupTableTest.java new file mode 100644 index 000000000000..95fb4a2b7385 --- /dev/null +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/join/LookupTableTest.java @@ -0,0 +1,484 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.runtime.operator.join; + +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.pinot.core.data.table.Key; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +import static org.testng.Assert.*; + + +public class LookupTableTest { + + @DataProvider(name = "lookupTableProvider") + public Object[][] provideLookupTables() { + return new Object[][]{ + {new IntLookupTable(), "IntLookupTable"}, + {new LongLookupTable(), "LongLookupTable"}, + {new FloatLookupTable(), "FloatLookupTable"}, + {new DoubleLookupTable(), "DoubleLookupTable"}, + {new ObjectLookupTable(), "ObjectLookupTable"} + }; + } + + @Test(dataProvider = "lookupTableProvider") + public void testEmptyTable(LookupTable table, String tableName) { + // Empty table should have unique keys by default + assertTrue(table.isKeysUnique(), tableName + " should start with unique keys"); + + // Lookup on empty table should return null + assertNull(table.lookup(getTestKey(table, 1)), tableName + " lookup should return null for empty table"); + + // ContainsKey on empty table should return false + assertFalse(table.containsKey(getTestKey(table, 1)), + tableName + " containsKey should return false for empty table"); + + // EntrySet should be empty + assertTrue(table.entrySet().isEmpty(), tableName + " entrySet should be empty"); + + // Finish should work on empty table + table.finish(); + } + + @Test(dataProvider = "lookupTableProvider") + public void testSingleRowUniqueKey(LookupTable table, String tableName) { + Object[] row1 = {"value1", 100, 1.5}; + Object key1 = getTestKey(table, 1); + + table.addRow(key1, row1); + table.finish(); + + // Should still be unique + assertTrue(table.isKeysUnique(), tableName + " should remain unique with single row"); + + // Should contain the key + assertTrue(table.containsKey(key1), tableName + " should contain the added key"); + + // Lookup should return the row directly (not as list) + Object result = table.lookup(key1); + assertNotNull(result, tableName + " lookup should not return null"); + assertTrue(result instanceof Object[], tableName + " lookup should return Object[] for unique keys"); + assertEquals(result, row1, tableName + " lookup should return the exact row"); + + // Should not contain other keys + assertFalse(table.containsKey(getTestKey(table, 2)), tableName + " should not contain non-added keys"); + + // EntrySet should have one entry + Set> entries = table.entrySet(); + assertEquals(entries.size(), 1, tableName + " entrySet should have exactly one entry"); + } + + @Test(dataProvider = "lookupTableProvider") + public void testMultipleRowsUniqueKeys(LookupTable table, String tableName) { + Object[] row1 = {"value1", 100, 1.5}; + Object[] row2 = {"value2", 200, 2.5}; + Object[] row3 = {"value3", 300, 3.5}; + + Object key1 = getTestKey(table, 1); + Object key2 = getTestKey(table, 2); + Object key3 = getTestKey(table, 3); + + table.addRow(key1, row1); + table.addRow(key2, row2); + table.addRow(key3, row3); + table.finish(); + + // Should remain unique + assertTrue(table.isKeysUnique(), tableName + " should remain unique with multiple unique keys"); + + // Should contain all keys + assertTrue(table.containsKey(key1), tableName + " should contain key1"); + assertTrue(table.containsKey(key2), tableName + " should contain key2"); + assertTrue(table.containsKey(key3), tableName + " should contain key3"); + + // Lookups should return correct rows + assertEquals(table.lookup(key1), row1, tableName + " should return correct row for key1"); + assertEquals(table.lookup(key2), row2, tableName + " should return correct row for key2"); + assertEquals(table.lookup(key3), row3, tableName + " should return correct row for key3"); + + // EntrySet should have three entries + assertEquals(table.entrySet().size(), 3, tableName + " entrySet should have three entries"); + } + + @Test(dataProvider = "lookupTableProvider") + public void testDuplicateKeys(LookupTable table, String tableName) { + Object[] row1 = {"value1", 100, 1.5}; + Object[] row2 = {"value2", 200, 2.5}; + Object[] row3 = {"value3", 300, 3.5}; + + Object key1 = getTestKey(table, 1); + + table.addRow(key1, row1); + table.addRow(key1, row2); // Duplicate key + table.addRow(key1, row3); // Another duplicate + table.finish(); + + // Should no longer be unique + assertFalse(table.isKeysUnique(), tableName + " should not be unique with duplicate keys"); + + // Should contain the key + assertTrue(table.containsKey(key1), tableName + " should contain the duplicated key"); + + // Lookup should return a list of rows + Object result = table.lookup(key1); + assertNotNull(result, tableName + " lookup should not return null for duplicate keys"); + assertTrue(result instanceof List, tableName + " lookup should return List for duplicate keys"); + + @SuppressWarnings("unchecked") + List resultList = (List) result; + assertEquals(resultList.size(), 3, tableName + " should return all three rows for duplicate key"); + + // Should contain all three rows in order + assertEquals(resultList.get(0), row1, tableName + " should return first row at index 0"); + assertEquals(resultList.get(1), row2, tableName + " should return second row at index 1"); + assertEquals(resultList.get(2), row3, tableName + " should return third row at index 2"); + + // EntrySet should have one entry (the key) with list value + Set> entries = table.entrySet(); + assertEquals(entries.size(), 1, tableName + " entrySet should have one entry for duplicate keys"); + } + + @Test(dataProvider = "lookupTableProvider") + public void testMixedUniqueAndDuplicateKeys(LookupTable table, String tableName) { + Object[] row1 = {"value1", 100, 1.5}; + Object[] row2 = {"value2", 200, 2.5}; + Object[] row3 = {"value3", 300, 3.5}; + Object[] row4 = {"value4", 400, 4.5}; + + Object key1 = getTestKey(table, 1); + Object key2 = getTestKey(table, 2); + + table.addRow(key1, row1); + table.addRow(key2, row2); // Unique key + table.addRow(key1, row3); // Duplicate of key1 + table.addRow(key1, row4); // Another duplicate of key1 + table.finish(); + + // Should not be unique + assertFalse(table.isKeysUnique(), tableName + " should not be unique with mixed keys"); + + // Should contain both keys + assertTrue(table.containsKey(key1), tableName + " should contain duplicated key"); + assertTrue(table.containsKey(key2), tableName + " should contain unique key"); + + // Key1 should return a list + Object result1 = table.lookup(key1); + assertTrue(result1 instanceof List, tableName + " should return List for duplicate key"); + @SuppressWarnings("unchecked") + List resultList1 = (List) result1; + assertEquals(resultList1.size(), 3, tableName + " should return three rows for duplicate key"); + + // Key2 should return a single-element list (due to finish() converting single rows to lists when not unique) + Object result2 = table.lookup(key2); + assertTrue(result2 instanceof List, tableName + " should return List for unique key when table has duplicates"); + @SuppressWarnings("unchecked") + List resultList2 = (List) result2; + assertEquals(resultList2.size(), 1, tableName + " should return one row for unique key"); + assertEquals(resultList2.get(0), row2, tableName + " should return correct row for unique key"); + + // EntrySet should have two entries + assertEquals(table.entrySet().size(), 2, tableName + " entrySet should have two entries"); + } + + @Test(dataProvider = "lookupTableProvider") + public void testNullKeyHandling(LookupTable table, String tableName) { + Object[] row1 = {"value1", 100, 1.5}; + Object[] row2 = {"value2", 200, 2.5}; + + Object key1 = getTestKey(table, 1); + + // Add normal row + table.addRow(key1, row1); + + // Add row with null key - should be ignored + table.addRow(null, row2); + + table.finish(); + + // Should remain unique since null key is ignored + assertTrue(table.isKeysUnique(), tableName + " should remain unique when null keys are ignored"); + + // Should contain the non-null key + assertTrue(table.containsKey(key1), tableName + " should contain non-null key"); + + // Should not contain null key + assertFalse(table.containsKey(null), tableName + " should not contain null key"); + + // Lookup with null should return null + assertNull(table.lookup(null), tableName + " lookup with null key should return null"); + + // Normal lookup should work + assertEquals(table.lookup(key1), row1, tableName + " should return correct row for non-null key"); + + // EntrySet should have only one entry (null key ignored) + assertEquals(table.entrySet().size(), 1, tableName + " entrySet should have one entry (null ignored)"); + } + + @Test(dataProvider = "lookupTableProvider") + public void testFinishBehaviorWithUniqueKeys(LookupTable table, String tableName) { + Object[] row1 = {"value1", 100, 1.5}; + Object[] row2 = {"value2", 200, 2.5}; + + Object key1 = getTestKey(table, 1); + Object key2 = getTestKey(table, 2); + + table.addRow(key1, row1); + table.addRow(key2, row2); + + // Before finish, should be unique + assertTrue(table.isKeysUnique(), tableName + " should be unique before finish"); + + table.finish(); + + // After finish with unique keys, should still be unique and return Object[] directly + assertTrue(table.isKeysUnique(), tableName + " should remain unique after finish"); + + Object result1 = table.lookup(key1); + Object result2 = table.lookup(key2); + + assertTrue(result1 instanceof Object[], tableName + " should return Object[] for unique keys after finish"); + assertTrue(result2 instanceof Object[], tableName + " should return Object[] for unique keys after finish"); + + assertEquals(result1, row1, tableName + " should return correct row after finish"); + assertEquals(result2, row2, tableName + " should return correct row after finish"); + } + + @Test(dataProvider = "lookupTableProvider") + public void testFinishBehaviorWithDuplicateKeys(LookupTable table, String tableName) { + Object[] row1 = {"value1", 100, 1.5}; + Object[] row2 = {"value2", 200, 2.5}; + Object[] row3 = {"value3", 300, 3.5}; + + Object key1 = getTestKey(table, 1); + Object key2 = getTestKey(table, 2); + + table.addRow(key1, row1); + table.addRow(key2, row2); + table.addRow(key1, row3); // Create duplicate + + // Before finish, should not be unique + assertFalse(table.isKeysUnique(), tableName + " should not be unique before finish with duplicates"); + + table.finish(); + + // After finish, should still not be unique + assertFalse(table.isKeysUnique(), tableName + " should remain non-unique after finish"); + + // All lookups should return Lists after finish with duplicates + Object result1 = table.lookup(key1); + Object result2 = table.lookup(key2); + + assertTrue(result1 instanceof List, tableName + " should return List for duplicate key after finish"); + assertTrue(result2 instanceof List, + tableName + " should return List for unique key after finish when table has duplicates"); + + @SuppressWarnings("unchecked") + List resultList1 = (List) result1; + @SuppressWarnings("unchecked") + List resultList2 = (List) result2; + + assertEquals(resultList1.size(), 2, tableName + " should return two rows for duplicate key"); + assertEquals(resultList2.size(), 1, tableName + " should return one row for unique key"); + } + + @Test + public void testObjectLookupTableWithComplexKeys() { + ObjectLookupTable table = new ObjectLookupTable(); + + // Test with composite keys (Key objects) + Key compositeKey1 = new Key(new Object[]{"string1", 100}); + Key compositeKey2 = new Key(new Object[]{"string2", 200}); + Key compositeKey3 = new Key(new Object[]{"string1", 100}); // Same as compositeKey1 + + Object[] row1 = {"value1", 1.0}; + Object[] row2 = {"value2", 2.0}; + Object[] row3 = {"value3", 3.0}; + + table.addRow(compositeKey1, row1); + table.addRow(compositeKey2, row2); + table.addRow(compositeKey3, row3); // Should be treated as duplicate of compositeKey1 + + table.finish(); + + // Should not be unique due to duplicate composite keys + assertFalse(table.isKeysUnique(), "ObjectLookupTable should not be unique with duplicate composite keys"); + + // Should contain the composite keys + assertTrue(table.containsKey(compositeKey1), "Should contain first composite key"); + assertTrue(table.containsKey(compositeKey2), "Should contain second composite key"); + assertTrue(table.containsKey(compositeKey3), "Should contain third composite key (same as first)"); + + // Lookup for duplicate key should return list + Object result1 = table.lookup(compositeKey1); + assertTrue(result1 instanceof List, "Should return List for duplicate composite key"); + @SuppressWarnings("unchecked") + List resultList1 = (List) result1; + assertEquals(resultList1.size(), 2, "Should return two rows for duplicate composite key"); + + // Lookup for unique key should return single-element list + Object result2 = table.lookup(compositeKey2); + assertTrue(result2 instanceof List, "Should return List for unique composite key when table has duplicates"); + @SuppressWarnings("unchecked") + List resultList2 = (List) result2; + assertEquals(resultList2.size(), 1, "Should return one row for unique composite key"); + } + + @Test + public void testObjectLookupTableWithStringKeys() { + ObjectLookupTable table = new ObjectLookupTable(); + + Object[] row1 = {"value1", 100}; + Object[] row2 = {"value2", 200}; + Object[] row3 = {"value3", 300}; + + table.addRow("key1", row1); + table.addRow("key2", row2); + table.addRow("key1", row3); // Duplicate string key + + table.finish(); + + assertFalse(table.isKeysUnique(), "ObjectLookupTable should not be unique with duplicate string keys"); + + assertTrue(table.containsKey("key1"), "Should contain key1"); + assertTrue(table.containsKey("key2"), "Should contain key2"); + + Object result1 = table.lookup("key1"); + assertTrue(result1 instanceof List, "Should return List for duplicate string key"); + @SuppressWarnings("unchecked") + List resultList1 = (List) result1; + assertEquals(resultList1.size(), 2, "Should return two rows for duplicate string key"); + } + + @Test + public void testIntLookupTableSpecifics() { + IntLookupTable table = new IntLookupTable(); + + Object[] row1 = {"value1", 100}; + Object[] row2 = {"value2", 200}; + + table.addRow(42, row1); + table.addRow(84, row2); + table.finish(); + + assertTrue(table.isKeysUnique(), "IntLookupTable should be unique"); + assertEquals(table.size(), 2, "IntLookupTable size should be 2"); + + assertTrue(table.containsKey(42), "Should contain int key 42"); + assertTrue(table.containsKey(84), "Should contain int key 84"); + assertFalse(table.containsKey(99), "Should not contain int key 99"); + + assertEquals(table.lookup(42), row1, "Should return correct row for int key 42"); + assertEquals(table.lookup(84), row2, "Should return correct row for int key 84"); + assertNull(table.lookup(99), "Should return null for non-existent int key"); + } + + @Test + public void testLongLookupTableSpecifics() { + LongLookupTable table = new LongLookupTable(); + + Object[] row1 = {"value1", 100}; + Object[] row2 = {"value2", 200}; + + table.addRow(42L, row1); + table.addRow(84L, row2); + table.finish(); + + assertTrue(table.isKeysUnique(), "LongLookupTable should be unique"); + assertEquals(table.size(), 2, "LongLookupTable size should be 2"); + + assertTrue(table.containsKey(42L), "Should contain long key 42L"); + assertTrue(table.containsKey(84L), "Should contain long key 84L"); + assertFalse(table.containsKey(99L), "Should not contain long key 99L"); + } + + @Test + public void testFloatLookupTableSpecifics() { + FloatLookupTable table = new FloatLookupTable(); + + Object[] row1 = {"value1", 100}; + Object[] row2 = {"value2", 200}; + + table.addRow(42.5f, row1); + table.addRow(84.7f, row2); + table.finish(); + + assertTrue(table.isKeysUnique(), "FloatLookupTable should be unique"); + assertEquals(table.size(), 2, "FloatLookupTable size should be 2"); + + assertTrue(table.containsKey(42.5f), "Should contain float key 42.5f"); + assertTrue(table.containsKey(84.7f), "Should contain float key 84.7f"); + assertFalse(table.containsKey(99.9f), "Should not contain float key 99.9f"); + } + + @Test + public void testDoubleLookupTableSpecifics() { + DoubleLookupTable table = new DoubleLookupTable(); + + Object[] row1 = {"value1", 100}; + Object[] row2 = {"value2", 200}; + + table.addRow(42.5, row1); + table.addRow(84.7, row2); + table.finish(); + + assertTrue(table.isKeysUnique(), "DoubleLookupTable should be unique"); + assertEquals(table.size(), 2, "DoubleLookupTable size should be 2"); + + assertTrue(table.containsKey(42.5), "Should contain double key 42.5"); + assertTrue(table.containsKey(84.7), "Should contain double key 84.7"); + assertFalse(table.containsKey(99.9), "Should not contain double key 99.9"); + } + + @Test + public void testObjectLookupTableSize() { + ObjectLookupTable table = new ObjectLookupTable(); + + Object[] row1 = {"value1", 100}; + Object[] row2 = {"value2", 200}; + + table.addRow("key1", row1); + table.addRow("key2", row2); + table.finish(); + + assertEquals(table.size(), 2, "ObjectLookupTable size should be 2"); + } + + /** + * Helper method to get appropriate test keys for different lookup table types + */ + private Object getTestKey(LookupTable table, int index) { + if (table instanceof IntLookupTable) { + return index; + } else if (table instanceof LongLookupTable) { + return (long) index; + } else if (table instanceof FloatLookupTable) { + return (float) index + 0.5f; + } else if (table instanceof DoubleLookupTable) { + return (double) index + 0.5; + } else if (table instanceof ObjectLookupTable) { + return "key" + index; + } else { + throw new IllegalArgumentException("Unknown lookup table type: " + table.getClass()); + } + } +} From ccd55239c5a8b0a9c0dcefadd90f68ad2ede30e7 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Wed, 23 Jul 2025 15:30:08 -0700 Subject: [PATCH 010/167] Bump commons-codec:commons-codec from 1.18.0 to 1.19.0 (#16413) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index d21f59fb6e7c..e5e7ac6dec47 100644 --- a/pom.xml +++ b/pom.xml @@ -210,7 +210,7 @@ 2.12.0 1.11.0 2.20.0 - 1.18.0 + 1.19.0 1.9.0 3.11.1 1.10.0 From da5109292984a6a15b6e4956769db8a4998a6eab Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Wed, 23 Jul 2025 15:30:23 -0700 Subject: [PATCH 011/167] Bump software.amazon.awssdk:bom from 2.32.5 to 2.32.6 (#16412) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index e5e7ac6dec47..ccccac7f5c3f 100644 --- a/pom.xml +++ b/pom.xml @@ -177,7 +177,7 @@ 0.15.1 0.4.7 4.2.2 - 2.32.5 + 2.32.6 1.2.36 1.22.0 2.14.0 From 0189f92523c7569e89604d74eb1c00e5a353bbf9 Mon Sep 17 00:00:00 2001 From: Chaitanya Deepthi <45308220+deepthi912@users.noreply.github.com> Date: Wed, 23 Jul 2025 17:30:41 -0700 Subject: [PATCH 012/167] Add a protected method for deleting snapshots to be extended (#16414) --- .../local/upsert/BasePartitionUpsertMetadataManager.java | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/BasePartitionUpsertMetadataManager.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/BasePartitionUpsertMetadataManager.java index ce6070b4ac56..b3da7e177cac 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/BasePartitionUpsertMetadataManager.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/BasePartitionUpsertMetadataManager.java @@ -358,7 +358,7 @@ protected void doAddSegment(ImmutableSegmentImpl segment) { } long startTimeMs = System.currentTimeMillis(); if (!_enableSnapshot) { - segment.deleteValidDocIdsSnapshot(); + deleteSnapshot(segment); } try (UpsertUtils.RecordInfoReader recordInfoReader = new UpsertUtils.RecordInfoReader(segment, _primaryKeyColumns, _comparisonColumns, _deleteRecordColumn)) { @@ -955,6 +955,10 @@ protected void doTakeSnapshot() { numConsumingSegments, System.currentTimeMillis() - startTimeMs); } + protected void deleteSnapshot(ImmutableSegmentImpl segment) { + segment.deleteValidDocIdsSnapshot(); + } + protected File getWatermarkFile() { return new File(_tableIndexDir, V1Constants.TTL_WATERMARK_TABLE_PARTITION + _partitionId); } From 1a3f2bd327813647833da165243fc212f489f904 Mon Sep 17 00:00:00 2001 From: Praveen Date: Wed, 23 Jul 2025 20:18:27 -0700 Subject: [PATCH 013/167] Emit segment download metrics for local (NFS/GPFS) file access (#16379) * local metrics * Update pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java Co-authored-by: Subbu Subramaniam * local metrics --------- Co-authored-by: Subbu Subramaniam --- ...tSegmentUploadDownloadRestletResource.java | 20 ++- .../api/resources/ResourceUtils.java | 4 +- .../api/resources/ResourceUtilsTest.java | 136 ++++++++++++++++++ 3 files changed, 154 insertions(+), 6 deletions(-) create mode 100644 pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/ResourceUtilsTest.java diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java index 76cd15db18dc..4cb0fcfde7f7 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java @@ -199,7 +199,19 @@ public Response downloadSegment( "Segment " + segmentName + " or table " + tableName + " not found in " + segmentFile.getAbsolutePath(), Response.Status.NOT_FOUND); } - builder.entity(segmentFile); + String rawTableName = TableNameBuilder.extractRawTableName(tableName); + long segmentSizeInBytes = segmentFile.length(); + long downloadStartTimeMs = System.currentTimeMillis(); + ResourceUtils.emitPreSegmentDownloadMetrics(_controllerMetrics, rawTableName, segmentSizeInBytes); + // Streaming the segment file directly from local FS to the output stream to ensure we can capture the metrics + builder.entity((StreamingOutput) output -> { + try { + Files.copy(segmentFile.toPath(), output); + } finally { + ResourceUtils.emitPostSegmentDownloadMetrics(_controllerMetrics, rawTableName, downloadStartTimeMs, + segmentSizeInBytes); + } + }); } else { URI remoteSegmentFileURI = URIUtils.getUri(dataDirURI.toString(), tableName, URIUtils.encode(segmentName)); PinotFS pinotFS = PinotFSFactory.create(dataDirURI.getScheme()); @@ -216,12 +228,12 @@ public Response downloadSegment( "Invalid segment name: %s", segmentName); String rawTableName = TableNameBuilder.extractRawTableName(tableName); // Emit metrics related to deep-store download operation - long deepStoreDownloadStartTimeMs = System.currentTimeMillis(); + long downloadStartTimeMs = System.currentTimeMillis(); long segmentSizeInBytes = segmentFile.length(); ResourceUtils.emitPreSegmentDownloadMetrics(_controllerMetrics, rawTableName, segmentSizeInBytes); pinotFS.copyToLocalFile(remoteSegmentFileURI, segmentFile); - ResourceUtils.emitPostSegmentDownloadMetrics(_controllerMetrics, rawTableName, - System.currentTimeMillis() - deepStoreDownloadStartTimeMs, segmentSizeInBytes); + ResourceUtils.emitPostSegmentDownloadMetrics(_controllerMetrics, rawTableName, downloadStartTimeMs, + segmentSizeInBytes); // Streaming in the tmp file and delete it afterward. builder.entity((StreamingOutput) output -> { diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/ResourceUtils.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/ResourceUtils.java index 3d6fce7448bf..4b799211e892 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/ResourceUtils.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/ResourceUtils.java @@ -113,9 +113,9 @@ public static void emitPostSegmentUploadMetrics(ControllerMetrics controllerMetr controllerMetrics.setValueOfGlobalGauge(ControllerGauge.DEEP_STORE_WRITE_OPS_IN_PROGRESS, writeCount); long segmentBytesUploading = _deepStoreWriteBytesInProgress.addAndGet(-segmentSizeInBytes); - controllerMetrics.setOrUpdateTableGauge(rawTableName, ControllerGauge.DEEP_STORE_WRITE_OPS_IN_PROGRESS, + controllerMetrics.setOrUpdateTableGauge(rawTableName, ControllerGauge.DEEP_STORE_WRITE_BYTES_IN_PROGRESS, segmentBytesUploading); - controllerMetrics.setValueOfGlobalGauge(ControllerGauge.DEEP_STORE_WRITE_OPS_IN_PROGRESS, segmentBytesUploading); + controllerMetrics.setValueOfGlobalGauge(ControllerGauge.DEEP_STORE_WRITE_BYTES_IN_PROGRESS, segmentBytesUploading); long durationMs = System.currentTimeMillis() - startTimeMs; controllerMetrics.addTimedTableValue(rawTableName, ControllerTimer.DEEP_STORE_SEGMENT_WRITE_TIME_MS, durationMs, diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/ResourceUtilsTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/ResourceUtilsTest.java new file mode 100644 index 000000000000..71eaa1a2eaf6 --- /dev/null +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/ResourceUtilsTest.java @@ -0,0 +1,136 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.controller.api.resources; + +import java.lang.reflect.Field; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; +import org.apache.pinot.common.metrics.ControllerGauge; +import org.apache.pinot.common.metrics.ControllerMeter; +import org.apache.pinot.common.metrics.ControllerMetrics; +import org.apache.pinot.common.metrics.ControllerTimer; +import org.mockito.ArgumentCaptor; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.testng.Assert.assertTrue; + +public class ResourceUtilsTest { + + private static final String RAW_TABLE = "myTable"; + private static final long SEGMENT_BYTES = 123L; + + private ControllerMetrics _metrics; + + /** Reset static Atomics inside ResourceUtils between tests so they don’t bleed values. */ + @BeforeMethod + public void setUp() throws Exception { + _metrics = mock(ControllerMetrics.class); + resetStaticCounter("_deepStoreWriteOpsInProgress"); + resetStaticCounter("_deepStoreWriteBytesInProgress"); + resetStaticCounter("_deepStoreReadOpsInProgress"); + resetStaticCounter("_deepStoreReadBytesInProgress"); + } + + @Test + public void testUploadMetricsFlow() throws InterruptedException { + long startTime = System.currentTimeMillis(); + + // This test verifies the flow of metrics emitted during a segment upload to deep store. + ResourceUtils.emitPreSegmentUploadMetrics(_metrics, RAW_TABLE, SEGMENT_BYTES); + + verify(_metrics).setOrUpdateTableGauge(RAW_TABLE, + ControllerGauge.DEEP_STORE_WRITE_OPS_IN_PROGRESS, 1L); + verify(_metrics).setOrUpdateTableGauge(RAW_TABLE, + ControllerGauge.DEEP_STORE_WRITE_BYTES_IN_PROGRESS, SEGMENT_BYTES); + + // Simulate upload latency + long simulatedDelay = 50L; + Thread.sleep(simulatedDelay); + + // Post-upload metrics + ResourceUtils.emitPostSegmentUploadMetrics(_metrics, RAW_TABLE, startTime, SEGMENT_BYTES); + + // ops / bytes gauges should be back to zero + verify(_metrics, atLeastOnce()).setOrUpdateTableGauge(RAW_TABLE, + ControllerGauge.DEEP_STORE_WRITE_OPS_IN_PROGRESS, 0L); + verify(_metrics, atLeastOnce()).setValueOfGlobalGauge(ControllerGauge.DEEP_STORE_WRITE_OPS_IN_PROGRESS, 0L); + verify(_metrics, atLeastOnce()).setOrUpdateTableGauge(RAW_TABLE, + ControllerGauge.DEEP_STORE_WRITE_BYTES_IN_PROGRESS, 0L); + verify(_metrics, atLeastOnce()).setValueOfGlobalGauge(ControllerGauge.DEEP_STORE_WRITE_BYTES_IN_PROGRESS, 0L); + + // timers & meters fire once + verify(_metrics).addTimedTableValue(eq(RAW_TABLE), + eq(ControllerTimer.DEEP_STORE_SEGMENT_WRITE_TIME_MS), anyLong(), eq(TimeUnit.MILLISECONDS)); + verify(_metrics).addMeteredTableValue(RAW_TABLE, + ControllerMeter.DEEP_STORE_WRITE_BYTES_COMPLETED, SEGMENT_BYTES); + // ArgumentCaptor to capture the long duration value that gets passed into the addTimedTableValue(...) method. + // This allows the test to inspect the actual argument used when the metric was recorded. + ArgumentCaptor durationCaptor = ArgumentCaptor.forClass(Long.class); + verify(_metrics).addTimedTableValue(eq(RAW_TABLE), + eq(ControllerTimer.DEEP_STORE_SEGMENT_WRITE_TIME_MS), durationCaptor.capture(), eq(TimeUnit.MILLISECONDS)); + assertTrue(durationCaptor.getValue() >= simulatedDelay, + "Expected write latency >= " + simulatedDelay + "ms but got " + durationCaptor.getValue()); + } + + @Test + public void testDownloadMetricsFlow() throws InterruptedException { + long startTime = System.currentTimeMillis(); + + ResourceUtils.emitPreSegmentDownloadMetrics(_metrics, RAW_TABLE, SEGMENT_BYTES); + verify(_metrics).setOrUpdateTableGauge(RAW_TABLE, + ControllerGauge.DEEP_STORE_READ_OPS_IN_PROGRESS, 1L); + verify(_metrics).setOrUpdateTableGauge(RAW_TABLE, + ControllerGauge.DEEP_STORE_READ_BYTES_IN_PROGRESS, SEGMENT_BYTES); + + // Simulate download latency + long simulatedDelay = 50L; + Thread.sleep(simulatedDelay); + // Post-download metrics + ResourceUtils.emitPostSegmentDownloadMetrics(_metrics, RAW_TABLE, startTime, SEGMENT_BYTES); + verify(_metrics, atLeastOnce()).setOrUpdateTableGauge(RAW_TABLE, + ControllerGauge.DEEP_STORE_READ_OPS_IN_PROGRESS, 0L); + verify(_metrics, atLeastOnce()).setOrUpdateTableGauge(RAW_TABLE, + ControllerGauge.DEEP_STORE_READ_BYTES_IN_PROGRESS, 0L); + + verify(_metrics).addTimedTableValue(eq(RAW_TABLE), + eq(ControllerTimer.DEEP_STORE_SEGMENT_READ_TIME_MS), anyLong(), eq(TimeUnit.MILLISECONDS)); + verify(_metrics).addMeteredTableValue(RAW_TABLE, + ControllerMeter.DEEP_STORE_READ_BYTES_COMPLETED, SEGMENT_BYTES); + // Capture and verify the read latency + // ArgumentCaptor to capture the long duration value that gets passed into the addTimedTableValue(...) method. + ArgumentCaptor durationCaptor = ArgumentCaptor.forClass(Long.class); + verify(_metrics).addTimedTableValue(eq(RAW_TABLE), eq(ControllerTimer.DEEP_STORE_SEGMENT_READ_TIME_MS), + durationCaptor.capture(), eq(TimeUnit.MILLISECONDS)); + assertTrue(durationCaptor.getValue() >= simulatedDelay, + "Expected read latency >= " + simulatedDelay + "ms but got " + durationCaptor.getValue()); + } + + /** resets the private static AtomicLongs inside ResourceUtils. */ + private static void resetStaticCounter(String fieldName) throws Exception { + Field declaredField = ResourceUtils.class.getDeclaredField(fieldName); + declaredField.setAccessible(true); + ((AtomicLong) declaredField.get(null)).set(0L); + } +} From 8316aa39a2e2ca85239b3d7249e43dbb235d92ed Mon Sep 17 00:00:00 2001 From: Gonzalo Ortiz Jaureguizar Date: Thu, 24 Jul 2025 09:04:46 +0200 Subject: [PATCH 014/167] udf test framework (#16258) --- .gitignore | 3 + .../common/function/FunctionRegistry.java | 54 +- .../common/function/PinotScalarFunction.java | 48 + .../function/TransformFunctionFactory.java | 5 + .../java/org/apache/pinot/core/udf/Udf.java | 169 ++ .../org/apache/pinot/core/udf/UdfExample.java | 201 ++ .../pinot/core/udf/UdfExampleBuilder.java | 296 +++ .../apache/pinot/core/udf/UdfParameter.java | 178 ++ .../apache/pinot/core/udf/UdfSignature.java | 73 + .../tests/BaseClusterIntegrationTest.java | 16 +- pinot-integration-tests/pom.xml | 5 + .../MultiStageEngineIntegrationTest.java | 2 +- .../pinot/integration/tests/udf/AvroSink.java | 191 ++ .../tests/udf/IntegrationUdfTestCluster.java | 227 ++ .../pinot/integration/tests/udf/UdfTest.java | 444 ++++ .../test/resources/udf-test-results/README.md | 35 + .../test/resources/udf-test-results/abs.yaml | 870 ++++++++ .../test/resources/udf-test-results/acos.yaml | 268 +++ .../resources/udf-test-results/adler32.yaml | 198 ++ .../udf-test-results/all-functions.yaml | 1952 +++++++++++++++++ .../test/resources/udf-test-results/and.yaml | 303 +++ .../resources/udf-test-results/array.yaml | 30 + .../udf-test-results/arrayConcatDouble.yaml | 339 +++ .../udf-test-results/arrayElementAtInt.yaml | 233 ++ .../resources/udf-test-results/arrayMax.yaml | 158 ++ .../udf-test-results/arrayaverage.yaml | 138 ++ .../resources/udf-test-results/degrees.yaml | 233 ++ .../resources/udf-test-results/isFalse.yaml | 144 ++ .../test/resources/udf-test-results/mult.yaml | 520 +++++ .../test/resources/udf-test-results/not.yaml | 163 ++ .../test/resources/udf-test-results/or.yaml | 373 ++++ .../test/resources/udf-test-results/plus.yaml | 520 +++++ .../udf-test-results/toDateTime.yaml | 128 ++ .../test/resources/udf-test-results/year.yaml | 163 ++ .../plugin/inputformat/avro/AvroUtils.java | 26 +- .../pinot/query/runtime/function/AbsUdf.java | 63 + .../pinot/query/runtime/function/AcosUdf.java | 73 + .../query/runtime/function/Adler32Udf.java | 66 + .../pinot/query/runtime/function/AndUdf.java | 82 + .../runtime/function/ArrayAverageUdf.java | 76 + .../function/ArrayConcatDoubleUdf.java | 71 + .../function/ArrayElementAtIntUdf.java | 65 + .../query/runtime/function/ArrayMaxUdf.java | 77 + .../query/runtime/function/ArrayUdf.java | 51 + .../query/runtime/function/DegreesUdf.java | 72 + .../query/runtime/function/IsFalseUdf.java | 75 + .../pinot/query/runtime/function/MultUdf.java | 81 + .../pinot/query/runtime/function/NotUdf.java | 82 + .../pinot/query/runtime/function/OrUdf.java | 82 + .../pinot/query/runtime/function/PlusUdf.java | 82 + .../query/runtime/function/ToDateTimeUdf.java | 67 + .../pinot/query/runtime/function/YearUdf.java | 71 + pinot-udf-test/pom.xml | 48 + .../udf/test/PinotFunctionEnvGenerator.java | 369 ++++ .../pinot/udf/test/ResultByExample.java | 233 ++ .../pinot/udf/test/UdfExampleResult.java | 97 + .../apache/pinot/udf/test/UdfReporter.java | 358 +++ .../apache/pinot/udf/test/UdfTestCluster.java | 70 + .../pinot/udf/test/UdfTestFramework.java | 288 +++ .../apache/pinot/udf/test/UdfTestResult.java | 220 ++ .../pinot/udf/test/UdfTestScenario.java | 63 + .../scenarios/AbstractUdfTestScenario.java | 125 ++ .../ExpressionTransformerTestScenario.java | 78 + .../IntermediateUdfTestScenario.java | 130 ++ .../scenarios/PredicateUdfTestScenario.java | 117 + .../TransformationUdfTestScenario.java | 85 + pom.xml | 6 + 67 files changed, 12218 insertions(+), 11 deletions(-) create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/udf/Udf.java create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/udf/UdfExample.java create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/udf/UdfExampleBuilder.java create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/udf/UdfParameter.java create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/udf/UdfSignature.java create mode 100644 pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/udf/AvroSink.java create mode 100644 pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/udf/IntegrationUdfTestCluster.java create mode 100644 pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/udf/UdfTest.java create mode 100644 pinot-integration-tests/src/test/resources/udf-test-results/README.md create mode 100644 pinot-integration-tests/src/test/resources/udf-test-results/abs.yaml create mode 100644 pinot-integration-tests/src/test/resources/udf-test-results/acos.yaml create mode 100644 pinot-integration-tests/src/test/resources/udf-test-results/adler32.yaml create mode 100644 pinot-integration-tests/src/test/resources/udf-test-results/all-functions.yaml create mode 100644 pinot-integration-tests/src/test/resources/udf-test-results/and.yaml create mode 100644 pinot-integration-tests/src/test/resources/udf-test-results/array.yaml create mode 100644 pinot-integration-tests/src/test/resources/udf-test-results/arrayConcatDouble.yaml create mode 100644 pinot-integration-tests/src/test/resources/udf-test-results/arrayElementAtInt.yaml create mode 100644 pinot-integration-tests/src/test/resources/udf-test-results/arrayMax.yaml create mode 100644 pinot-integration-tests/src/test/resources/udf-test-results/arrayaverage.yaml create mode 100644 pinot-integration-tests/src/test/resources/udf-test-results/degrees.yaml create mode 100644 pinot-integration-tests/src/test/resources/udf-test-results/isFalse.yaml create mode 100644 pinot-integration-tests/src/test/resources/udf-test-results/mult.yaml create mode 100644 pinot-integration-tests/src/test/resources/udf-test-results/not.yaml create mode 100644 pinot-integration-tests/src/test/resources/udf-test-results/or.yaml create mode 100644 pinot-integration-tests/src/test/resources/udf-test-results/plus.yaml create mode 100644 pinot-integration-tests/src/test/resources/udf-test-results/toDateTime.yaml create mode 100644 pinot-integration-tests/src/test/resources/udf-test-results/year.yaml create mode 100644 pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/AbsUdf.java create mode 100644 pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/AcosUdf.java create mode 100644 pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/Adler32Udf.java create mode 100644 pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/AndUdf.java create mode 100644 pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/ArrayAverageUdf.java create mode 100644 pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/ArrayConcatDoubleUdf.java create mode 100644 pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/ArrayElementAtIntUdf.java create mode 100644 pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/ArrayMaxUdf.java create mode 100644 pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/ArrayUdf.java create mode 100644 pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/DegreesUdf.java create mode 100644 pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/IsFalseUdf.java create mode 100644 pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/MultUdf.java create mode 100644 pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/NotUdf.java create mode 100644 pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/OrUdf.java create mode 100644 pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/PlusUdf.java create mode 100644 pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/ToDateTimeUdf.java create mode 100644 pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/YearUdf.java create mode 100644 pinot-udf-test/pom.xml create mode 100644 pinot-udf-test/src/main/java/org/apache/pinot/udf/test/PinotFunctionEnvGenerator.java create mode 100644 pinot-udf-test/src/main/java/org/apache/pinot/udf/test/ResultByExample.java create mode 100644 pinot-udf-test/src/main/java/org/apache/pinot/udf/test/UdfExampleResult.java create mode 100644 pinot-udf-test/src/main/java/org/apache/pinot/udf/test/UdfReporter.java create mode 100644 pinot-udf-test/src/main/java/org/apache/pinot/udf/test/UdfTestCluster.java create mode 100644 pinot-udf-test/src/main/java/org/apache/pinot/udf/test/UdfTestFramework.java create mode 100644 pinot-udf-test/src/main/java/org/apache/pinot/udf/test/UdfTestResult.java create mode 100644 pinot-udf-test/src/main/java/org/apache/pinot/udf/test/UdfTestScenario.java create mode 100644 pinot-udf-test/src/main/java/org/apache/pinot/udf/test/scenarios/AbstractUdfTestScenario.java create mode 100644 pinot-udf-test/src/main/java/org/apache/pinot/udf/test/scenarios/ExpressionTransformerTestScenario.java create mode 100644 pinot-udf-test/src/main/java/org/apache/pinot/udf/test/scenarios/IntermediateUdfTestScenario.java create mode 100644 pinot-udf-test/src/main/java/org/apache/pinot/udf/test/scenarios/PredicateUdfTestScenario.java create mode 100644 pinot-udf-test/src/main/java/org/apache/pinot/udf/test/scenarios/TransformationUdfTestScenario.java diff --git a/.gitignore b/.gitignore index bbd3d3465637..5fb3db21a063 100644 --- a/.gitignore +++ b/.gitignore @@ -58,3 +58,6 @@ kubernetes/helm/**/Chart.lock #Develocity .mvn/.gradle-enterprise/ .mvn/.develocity/ + +pinot-integration-tests/src/test/resources/udf-test-results/*.md +!/pinot-integration-tests/src/test/resources/udf-test-results/README.md \ No newline at end of file diff --git a/pinot-common/src/main/java/org/apache/pinot/common/function/FunctionRegistry.java b/pinot-common/src/main/java/org/apache/pinot/common/function/FunctionRegistry.java index 3e0da5733d5d..75f5dc1eb8db 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/function/FunctionRegistry.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/function/FunctionRegistry.java @@ -22,10 +22,13 @@ import java.lang.reflect.Method; import java.lang.reflect.Modifier; import java.util.Arrays; +import java.util.Collections; import java.util.HashMap; import java.util.HashSet; +import java.util.List; import java.util.Map; import java.util.Set; +import java.util.stream.Collectors; import javax.annotation.Nullable; import org.apache.calcite.rel.type.RelDataType; import org.apache.calcite.rel.type.RelDataTypeFactory; @@ -76,11 +79,13 @@ private FunctionRegistry() { // Key is canonical name public static final Map FUNCTION_MAP; - private static final int VAR_ARG_KEY = -1; + public static final int VAR_ARG_KEY = -1; static { long startTimeMs = System.currentTimeMillis(); + // TODO: Register functions for UDFs + // Register ScalarFunction classes Map functionMap = new HashMap<>(); Set> classes = @@ -184,6 +189,10 @@ private static void register(String canonicalName, FunctionInfo functionInfo, in numArguments == VAR_ARG_KEY ? "variable" : numArguments); } + public static Map getFunctions() { + return Collections.unmodifiableMap(FUNCTION_MAP); + } + /** * Returns {@code true} if the given canonical name is registered, {@code false} otherwise. * @@ -241,29 +250,43 @@ public static String canonicalize(String name) { } public static class ArgumentCountBasedScalarFunction implements PinotScalarFunction { - private final String _name; + private final String _mainName; + private final Set _names; private final Map _functionInfoMap; - private ArgumentCountBasedScalarFunction(String name, Map functionInfoMap) { - _name = name; + public ArgumentCountBasedScalarFunction(String name, Map functionInfoMap) { + this(List.of(name), functionInfoMap); + } + + public ArgumentCountBasedScalarFunction(List names, Map functionInfoMap) { + Preconditions.checkArgument(!names.isEmpty(), "At least one name is required"); + _mainName = FunctionRegistry.canonicalize(names.get(0)); + _names = names.stream() + .map(FunctionRegistry::canonicalize) + .collect(Collectors.toSet()); _functionInfoMap = functionInfoMap; } @Override public String getName() { - return _name; + return _mainName; + } + + @Override + public Set getNames() { + return _names; } @Override public PinotSqlFunction toPinotSqlFunction() { - return new PinotSqlFunction(_name, getReturnTypeInference(), getOperandTypeChecker()); + return new PinotSqlFunction(_mainName, getReturnTypeInference(), getOperandTypeChecker()); } private SqlReturnTypeInference getReturnTypeInference() { return opBinding -> { int numArguments = opBinding.getOperandCount(); FunctionInfo functionInfo = getFunctionInfo(numArguments); - Preconditions.checkState(functionInfo != null, "Failed to find function: %s with %s arguments", _name, + Preconditions.checkState(functionInfo != null, "Failed to find function: %s with %s arguments", _mainName, numArguments); RelDataTypeFactory typeFactory = opBinding.getTypeFactory(); Method method = functionInfo.getMethod(); @@ -325,5 +348,22 @@ public FunctionInfo getFunctionInfo(int numArguments) { FunctionInfo functionInfo = _functionInfoMap.get(numArguments); return functionInfo != null ? functionInfo : _functionInfoMap.get(VAR_ARG_KEY); } + + @Override + public String getScalarFunctionId() { + if (_functionInfoMap.size() == 1) { + String singleFunInfo = printFunctionInfo(_functionInfoMap.values().iterator().next()); + return "ArgumentCountBasedScalarFunction{" + singleFunInfo + "}"; + } + String mapDescription = _functionInfoMap.entrySet().stream() + .map(pair -> pair.getKey() + ": " + printFunctionInfo(pair.getValue())) + .collect(Collectors.joining(", ", "[", "]")); + return "ArgumentCountBasedScalarFunction{" + mapDescription + "}"; + } + + private String printFunctionInfo(FunctionInfo functionInfo) { + Method method = functionInfo.getMethod(); + return method.getDeclaringClass().getTypeName() + '.' + method.getName(); + } } } diff --git a/pinot-common/src/main/java/org/apache/pinot/common/function/PinotScalarFunction.java b/pinot-common/src/main/java/org/apache/pinot/common/function/PinotScalarFunction.java index 6a2ed5e626ab..1074d6c4f3e3 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/function/PinotScalarFunction.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/function/PinotScalarFunction.java @@ -18,6 +18,11 @@ */ package org.apache.pinot.common.function; +import java.lang.reflect.Method; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.Set; import javax.annotation.Nullable; import org.apache.pinot.common.function.sql.PinotSqlFunction; import org.apache.pinot.common.utils.DataSchema.ColumnDataType; @@ -35,6 +40,36 @@ public interface PinotScalarFunction { */ String getName(); + /** + * Returns the set of names of the function, including the primary name returned by {@link #getName()}. + * + * The value of this function may be used to register the function in the OperatorTable, although the names included + * an optional {@link ScalarFunction}} annotation higher priority. + */ + default Set getNames() { + return Set.of(getName()); + } + + /** + * Returns an identifier for the scalar function, which is used to identify the actual PinotScalarFunction class that + * is registered in the FunctionRegistry. + * + * This was originally created to facilitate the debug process, so it is possible to identify the class that + * implements a given scalar function. This is for example used to generate the all-functions.yml file, which lists + * all the scalar and transform functions available in Pinot in order to test UDF implementations. + * + * See for example {@link FunctionRegistry.ArgumentCountBasedScalarFunction} which overrides this method to also + * include the different FunctionInfo for the different argument counts. + * + * It is important that this method returns a stable identifier. This means it should not change unless we change the + * class. For example, this should not be a call to {@link java.util.Objects#toIdentityString(Object)}, given it will + * include the hash code of the object id, which will be different for each instance of the object. Instead, it should + * depend on the class. + */ + default String getScalarFunctionId() { + return getClass().getCanonicalName(); + } + /** * Returns the corresponding {@link PinotSqlFunction} to be registered into the OperatorTable, or {@code null} if it * doesn't need to be registered (e.g. standard SqlFunction). @@ -57,4 +92,17 @@ default FunctionInfo getFunctionInfo(ColumnDataType[] argumentTypes) { */ @Nullable FunctionInfo getFunctionInfo(int numArguments); + + static PinotScalarFunction fromMethod(Method method, boolean isVarArg, boolean supportNullArgs, + @Nullable String... names) { + int numArguments = isVarArg ? FunctionRegistry.VAR_ARG_KEY : method.getParameterCount(); + FunctionInfo functionInfo = new FunctionInfo(method, method.getDeclaringClass(), supportNullArgs); + Map functionInfoMap = Map.of(numArguments, functionInfo); + + List nameList = names != null && names.length > 0 + ? Arrays.asList(names) + : List.of(method.getName()); + + return new FunctionRegistry.ArgumentCountBasedScalarFunction(nameList, functionInfoMap); + } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java index b2521e768146..f25f51d8d33b 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/TransformFunctionFactory.java @@ -20,6 +20,7 @@ import com.google.common.annotations.VisibleForTesting; import java.util.ArrayList; +import java.util.Collections; import java.util.EnumMap; import java.util.HashMap; import java.util.List; @@ -405,4 +406,8 @@ public static TransformFunction getNullHandlingEnabled(ExpressionContext express public static String canonicalize(String functionName) { return StringUtils.remove(functionName, '_').toLowerCase(); } + + public static Map> getAllFunctions() { + return Collections.unmodifiableMap(TRANSFORM_FUNCTION_MAP); + } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/udf/Udf.java b/pinot-core/src/main/java/org/apache/pinot/core/udf/Udf.java new file mode 100644 index 000000000000..fd37205b7ea4 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/udf/Udf.java @@ -0,0 +1,169 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.core.udf; + +import java.lang.reflect.Method; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; +import javax.annotation.Nullable; +import org.apache.arrow.util.Preconditions; +import org.apache.commons.lang3.tuple.Pair; +import org.apache.pinot.common.function.FunctionRegistry; +import org.apache.pinot.common.function.PinotScalarFunction; +import org.apache.pinot.common.function.TransformFunctionType; +import org.apache.pinot.core.operator.transform.function.TransformFunction; +import org.apache.pinot.spi.annotations.ScalarFunction; + + +/// The Udf interface represents a User Defined Function (UDF) in Pinot. +/// +/// In Pinot UDFs can be either [@ScalarFunction][org.apache.pinot.spi.annotations.ScalarFunction] or +/// [TransformFunction][org.apache.pinot.core.operator.transform.function.TransformFunction]. +/// The first are row based (are called once per row) and the second are block based (are called once per block), which +/// makes them more efficient for large datasets. +/// +/// These functions can be used in different parts of the Pinot query processing pipeline. For example, +/// TransformFunctions are when ProjectPlanNodes in SSE are materialized into TransformOpeartors, while scalar functions +/// are used mostly everywhere else, such as in filter expressions or even project nodes in MSE. +/// But although a ScalarFunction can be wrapped into a TransformFunction using ScalarTransformFunctionWrapper, +/// TransformFunctions cannot be used in places where ScalarFunctions are expected. +/// +/// Therefore in order to add a new function, one should always implement the ScalarFunction interface, and if +/// TransformFunction is needed, it can be implemented as a wrapper around the ScalarFunction. But this was not +/// automatically enforced by the APIs. +/// +/// This is why the Udf function was introduced. Udf interfaces should be the actual way to register an UDF in Pinot. +/// This interface is used to provide a unified way to describe UDFs, including their main function name, +/// description, examples and in future it could be used to register functions in TransformFunctionFactory and +/// FunctionRegistry (which is the one used to look for scalar functions). +/// +/// The examples are used to provide a set of examples for the function, which can be used in documentation or testing. +public abstract class Udf { + + /// The main function name of the UDF. + /// + /// This is treated as an ID, which means that on a single Pinot process there should be only one UDF with a given + /// main function name. + public abstract String getMainName(); + + /// Returns the main function name of the UDF, canonicalized as defined in [FunctionRegistry#canonicalize]. + /// This is used to ensure that the function name is in a consistent format, which is important for + /// function registration, lookup and reporting. + public String getMainCanonicalName() { + return FunctionRegistry.canonicalize(getMainName()); + } + + /// A set with all names of the functions that this UDF can be called with, including the main name. + /// + /// This is used to support different aliases for the same function, so that users can call the function. + public Set getAllNames() { + return Set.of(getMainName()); + } + + /// Returns a set with all names of the functions that this UDF can be called with, including the main name, + /// canonicalized as defined in [FunctionRegistry#canonicalize]. + /// + /// This is used to ensure that the function names are in a consistent format, which is important for + /// function registration, lookup and reporting. + public Set getAllCanonicalNames() { + return getAllNames().stream() + .map(FunctionRegistry::canonicalize) + .collect(Collectors.toSet()); + } + + /// A description of the UDF, which should be used in documentation or for debugging purposes. + /// + /// The description should be a human-readable text in markdown format that explains what the function does, + // language=markdown + public abstract String getDescription(); + + /// Returns the text that should be used in a SQL query to call the function. + /// + /// This is used to generate the SQL call for the function in test cases or documentation. + /// + /// @param name the name to be used. It should be one of the names returned by getAllFunctionNames(). + /// @param sqlArgValues the values of the arguments to be used in the SQL call. They can be either field references or + /// literal values, depending on the test case. + public String asSqlCall(String name, List sqlArgValues) { + return name + "(" + String.join(", ", sqlArgValues) + ")"; + } + + /// Returns the examples for this Udf. + /// + /// As UDFs can have multiple signatures, the examples are grouped by them. + /// + /// It is recommended to use [UdfExampleBuilder] in order to build the examples for the Udf. + public abstract Map> getExamples(); + + /// The pair of function type and transform function that implements this UDF. + /// + /// Unstable API: Transform functions still use an old model of registration using a model that is not polymorphic + @Nullable + public Pair> getTransformFunction() { + return null; + } + + /// Returns the ScalarFunctions that implement this UDF, if any. + /// + /// Ideally all UDFs should have a corresponding ScalarFunction. Otherwise Pinot won't be able to use them in some + /// scenarios. This should be enforced for all new UDFs, but existing UDFs might not have a corresponding + /// ScalarFunction, in which case they can return null. + @Nullable + public abstract PinotScalarFunction getScalarFunction(); + + @Override + public String toString() { + return getMainName(); + } + + public static abstract class FromAnnotatedMethod extends Udf { + private final Method _method; + private final ScalarFunction _annotation; + + public FromAnnotatedMethod(Method method) { + _method = method; + _annotation = Preconditions.checkNotNull(method.getAnnotation(ScalarFunction.class), + "Method %s is not annotated with @ScalarFunction", method); + } + + @Override + public Set getAllNames() { + if (_annotation.names().length > 0) { + return Set.of(_annotation.names()); + } + return Set.of(_method.getName()); + } + + @Override + public String getMainName() { + if (_annotation.names().length > 0) { + return _annotation.names()[0]; + } + return _method.getName(); + } + + @Override + public PinotScalarFunction getScalarFunction() { + return PinotScalarFunction.fromMethod(_method, _annotation.isVarArg(), + _annotation.nullableParameters(), _annotation.names()); + } + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/udf/UdfExample.java b/pinot-core/src/main/java/org/apache/pinot/core/udf/UdfExample.java new file mode 100644 index 000000000000..8fc40bce39d4 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/udf/UdfExample.java @@ -0,0 +1,201 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.core.udf; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Objects; + +/// The inputs and the expected result for a given [UDF][Udf] call. +/// +/// These objects are usually created by the [#create(String, Object...)] factory method and decorated when needed +/// calling [#withoutNull(Object)]. +public abstract class UdfExample { + /// A descriptive ID of what this example is showing. + public abstract String getId(); + + /// The input values. + /// + /// Notice that values should be the same used to call the function and should match the types indicated by the + /// signature associated with this example. For example, in the case of plus udf, args could be doubles 1.0 and 2.0 or + /// ints 1 and 2, but not double 1.0 and int 2. That sum is possible at SQL level, but in that case the different + /// engines apply automatic casting to make sure the udf is called with either two ints or two doubles. + public abstract List getInputValues(); + + /// The result of the example. It may be different depending on whether null handling is enabled or not. + public abstract Object getResult(NullHandling nullHandling); + + /// Creates a new UdfExample with the given name and some values. + /// + /// The values array is processed assuming that the last value is the expected result, while all the previous values + /// are the inputs to the UDF. + /// + /// The returned UdfExample will always return the same result regardless of whether null handling is enabled or not. + /// In case you want to create a UdfExample that returns different results depending on null handling, you should + /// call [#withoutNull(Object)] on the returned UdfExample and pass the expected result when null handling is + /// disabled. + public static UdfExample create(String name, Object... values) { + int inputsSize = values.length - 1; + ArrayList inputs = new ArrayList<>(inputsSize); + inputs.addAll(Arrays.asList(values).subList(0, inputsSize)); + Object expectedResult = values[inputsSize]; + return new Default(name, inputs, expectedResult); + } + + /// Returns a new UdfExample that will use the given result when null handling is disabled. + /// + /// Remember that normally the result is the same regardless of null handling, but in some cases it may be different. + public UdfExample withoutNull(Object result) { + return new WithoutNullHandling(this, result); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof UdfExample)) { + return false; + } + UdfExample that = (UdfExample) o; + if (!getId().equals(that.getId())) { + return false; + } + List myInputs = getInputValues(); + List otherInputs = that.getInputValues(); + if (myInputs.size() != otherInputs.size()) { + return false; + } + for (int i = 0; i < myInputs.size(); i++) { + if (!equalValues(myInputs.get(i), otherInputs.get(i))) { + return false; + } + } + return equalValues(getResult(NullHandling.DISABLED), that.getResult(NullHandling.DISABLED)) + && equalValues(getResult(NullHandling.ENABLED), that.getResult(NullHandling.ENABLED)); + } + + private boolean equalValues(Object result1, Object result2) { + if (result1 == null && result2 == null) { + return true; + } + if (result1 == null || result2 == null) { + return false; + } + if (result1.getClass().isArray() && result2.getClass().isArray()) { + if (!result1.getClass().equals(result2.getClass())) { + return false; + } + if (result1.getClass().equals(byte[].class)) { + return Arrays.equals((byte[]) result1, (byte[]) result2); + } + return Arrays.equals((Object[]) result1, (Object[]) result2); + } + return Objects.equals(result1, result2); + } + + @Override + public int hashCode() { + return Objects.hashCode(getId()); + } + + public static class Default extends UdfExample { + private final String _id; + private final List _inputValues; + private final Object _expectedResult; + + private Default(String id, List inputValues, Object expectedResult) { + _id = id; + _inputValues = inputValues; + _expectedResult = expectedResult; + } + + @Override + public String getId() { + return _id; + } + + @Override + public List getInputValues() { + return _inputValues; + } + + @Override + public Object getResult(NullHandling nullHandling) { + return _expectedResult; + } + } + + /// A decorator for UdfExample that allows to return a different result when null handling is disabled. + public static class WithoutNullHandling extends UdfExample { + private final UdfExample _base; + private final Object _result; + + private WithoutNullHandling(UdfExample base, Object result) { + _base = base; + _result = result; + } + + @Override + public List getInputValues() { + return _base.getInputValues(); + } + + @Override + public String getId() { + return _base.getId(); + } + + @Override + public Object getResult(NullHandling nullHandling) { + return nullHandling == NullHandling.DISABLED ? _result : _base.getResult(NullHandling.ENABLED); + } + + /// Returns a new UdfExample that uses the given result when null handling is disabled instead of the one provided + /// by the receiver object. + /// + /// This is similar to calling [#withoutNull(Object)] on the decorated object. + /// Calling this method doesn't seem to make much sense, but it is provided for consistency with the decorated + /// class. + @Override + public UdfExample withoutNull(Object result) { + return _base.withoutNull(result); + } + + @Override + public boolean equals(Object o) { + if (!(o instanceof WithoutNullHandling)) { + return false; + } + WithoutNullHandling that = (WithoutNullHandling) o; + return Objects.equals(_base, that._base) && Objects.equals(_result, that._result); + } + + @Override + public int hashCode() { + return Objects.hash(_base, _result); + } + } + + public enum NullHandling { + ENABLED, + DISABLED + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/udf/UdfExampleBuilder.java b/pinot-core/src/main/java/org/apache/pinot/core/udf/UdfExampleBuilder.java new file mode 100644 index 000000000000..9c39906fed18 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/udf/UdfExampleBuilder.java @@ -0,0 +1,296 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.core.udf; + +import com.google.common.collect.Maps; +import java.math.BigDecimal; +import java.util.Arrays; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; +import org.apache.arrow.util.Preconditions; +import org.apache.pinot.spi.data.FieldSpec; + +/// A builder for generating test cases for [UDFs][Udf] (User Defined Functions). +/// +/// It allows to create test cases for a specific UDF signature, or for endomorphism functions that operate on +/// numeric types. +public interface UdfExampleBuilder { + Map> generateExamples(); + + /// Starts a builder for the given signature. + static SingleBuilder forSignature(UdfSignature signature) { + return new SingleBuilder(signature); + } + + /// Starts an endomorphism numeric builder for the given arity. + /// + /// This builder is used for functions that take a bunch of numeric values (ie ints) and return the same numeric type. + /// Most mathematical functions are endomorphism functions, like `abs`, `plus`, `mult` etc. + /// + /// These samples can always be generated using [#forSignature], but that would require to write the same test cases + /// for each numeric type (int, double, big_decimal, etc), which is not very convenient. + /// Given each numeric type has its own properties (specially the precision), it is still recommended to write some + /// specific samples for each numeric type, but this builder can be used to generate the most common cases. + static EndomorphismNumericTypesBuilder forEndomorphismNumeric(int arity) { + return new EndomorphismNumericTypesBuilder(arity); + } + + class SingleBuilder { + private final UdfSignature _signature; + private final Set _examples = new HashSet<>(); + + protected SingleBuilder(UdfSignature signature) { + _signature = signature; + } + + /// Adds a new example for the current signature. + /// + /// It is recommended to use the `addCase(String, Object...)` method instead, as it is more convenient. + public SingleBuilder addExample(UdfExample example) { + Preconditions.checkArgument(example.getInputValues().size() == _signature.getArity(), + "Expected %s input values for signature %s, but got %s", + _signature.getArity(), _signature, example.getInputValues().size()); + _examples.add(example); + return this; + } + + /// Adds a new example for the current signature with the given name and values. + /// + /// This is equivalent to call `addCase(UdfExample.create(name, values))`. + /// The created example will use the latest value as the expected result of the function. + /// Remember that UdfExamples support up to 2 return values, one for when null handling is enabled and one + /// for when null handling is disabled. + /// This method assumes that both are going to be the same. In case they are not, create the example using the + /// [UdfExample#create(String, Object...)] method, decorate it with [UdfExample#withoutNull(Object)] method + /// and finally call [SingleBuilder#addExample(UdfExample)]. + /// + /// @param name the name of the example. It should be unique for the given signature. + /// @param values the values for the example. The last value is the expected result of the function while all the + /// other values are the input values. This means that the length of the values array must be + /// the arity of the signature plus 1. + public SingleBuilder addExample(String name, Object... values) { + UdfExample example = UdfExample.create(name, values); + return addExample(example); + } + + /// Adds a new example, which may have a different signature than the current one. + /// This is used to combine multiple builders in order to create a compound map of examples. + public MultiBuilder and(UdfExampleBuilder other) { + MultiBuilder multiBuilder = new MultiBuilder(); + multiBuilder.and(build()); + multiBuilder.and(other); + return multiBuilder; + } + + /// Creates the UdfExampleBuilder that generates the examples for the current signature. + public UdfExampleBuilder build() { + return () -> Map.of(_signature, _examples); + } + } + + /// A builder for generating cases for endomorphism functions that operate on numeric types. + /// + /// It means that it takes a bunch of Number as arguments and returns the same Number. For example, if it takes + /// two ints, it returns an int. If it takes two longs, it returns a long, etc. + class EndomorphismNumericTypesBuilder { + private final int _arity; + private final Set _cases = new HashSet<>(); + + protected EndomorphismNumericTypesBuilder(int arity) { + _arity = arity; + } + + /// Like [SingleBuilder#addExample], but for endomorphism functions. + public EndomorphismNumericTypesBuilder addExample(String name, Number... values) { + Preconditions.checkArgument(values.length == _arity + 1, + "Expected %s values, but got %s", _arity + 1, values.length); + UdfExample example = UdfExample.create(name, values); + return addExample(example); + } + + /// Like [SingleBuilder#addExampleWithoutNull], but for endomorphism functions. + public EndomorphismNumericTypesBuilder addExampleWithoutNull(String name, Number... values) { + Preconditions.checkArgument(values.length == _arity + 2, + "Expected %s values, but got %s", _arity + 2, values.length); + // Copy the first _arity values and add the value returned when null handling is disabled. + Number[] nullableValues = Arrays.copyOf(values, _arity + 1); + UdfExample example = UdfExample.create(name, (Object[]) nullableValues) + // Ad the expected result when null handling is disabled. This value is at the end of the values array. + // and values.length - 1 == _arity + 1. + .withoutNull(values[_arity + 1]); + + return addExample(example); + } + + /// Like [SingleBuilder#addExample], but for endomorphism functions. + public EndomorphismNumericTypesBuilder addExample(UdfExample example) { + Preconditions.checkArgument(example.getInputValues().size() == _arity, + "Expected %s input values for endomorphism function with arity %s, but got %s", + _arity, _arity, example.getInputValues().size()); + _cases.add(example); + return this; + } + + public MultiBuilder and(UdfExampleBuilder other) { + MultiBuilder multiBuilder = new MultiBuilder(); + multiBuilder.and(build()); + multiBuilder.and(other); + return multiBuilder; + } + + /// Builds the UdfExampleBuilder that generates the examples for endomorphism functions. + public UdfExampleBuilder build() { + return () -> { + Set numericTypes = Set.of( + FieldSpec.DataType.INT, + FieldSpec.DataType.LONG, + FieldSpec.DataType.FLOAT, + FieldSpec.DataType.DOUBLE, + FieldSpec.DataType.BIG_DECIMAL + ); + Map> cases = Maps.newHashMapWithExpectedSize(numericTypes.size()); + for (FieldSpec.DataType type : numericTypes) { + UdfParameter[] paramsAndResult = new UdfParameter[_arity + 1]; + for (int i = 0; i < _arity; i++) { + paramsAndResult[i] = UdfParameter.of("arg" + i, type); + } + paramsAndResult[_arity] = UdfParameter.result(type); + UdfSignature signature = UdfSignature.of(paramsAndResult); + Set newCases = _cases.stream() + .map(testCase -> new EndomorphismNumericTypesBuilder.NumericPinotUdfExample(testCase, type)) + .collect(Collectors.toSet()); + cases.put(signature, newCases); + } + return cases; + }; + } + + private static class NumericPinotUdfExample extends UdfExample { + private final UdfExample _base; + private final FieldSpec.DataType _type; + + public NumericPinotUdfExample(UdfExample base, FieldSpec.DataType type) { + Preconditions.checkArgument(base.getResult(NullHandling.ENABLED) instanceof Number + || base.getResult(NullHandling.ENABLED) == null, + "Base test case must return a Number type for numeric endomorphism functions"); + Preconditions.checkArgument(base.getResult(NullHandling.DISABLED) instanceof Number + || base.getResult(NullHandling.DISABLED) == null, + "Base test case must return a Number type for numeric endomorphism functions"); + if (base.getInputValues().stream() + .anyMatch(value -> value != null && !(value instanceof Number))) { + throw new IllegalStateException( + "Base test case must have all input values as Number type for numeric endomorphism functions, got: " + + base.getInputValues().stream() + .map(value -> value + " (of type " + value.getClass().getName() + ")") + .collect(Collectors.joining(", ")) + ); + } + Preconditions.checkArgument(type == FieldSpec.DataType.INT + || type == FieldSpec.DataType.LONG + || type == FieldSpec.DataType.FLOAT + || type == FieldSpec.DataType.DOUBLE + || type == FieldSpec.DataType.BIG_DECIMAL, + "Type %s is not a numeric for numeric type", type); + _base = base; + _type = type; + } + + @Override + public String getId() { + return _base.getId() + "_" + _type.name().toLowerCase(Locale.US); + } + + @Override + public List getInputValues() { + return _base.getInputValues().stream() + .map(value -> { + if (value instanceof Number) { + switch (_type) { + case INT: + return ((Number) value).intValue(); + case LONG: + return ((Number) value).longValue(); + case FLOAT: + return ((Number) value).floatValue(); + case DOUBLE: + return ((Number) value).doubleValue(); + case BIG_DECIMAL: + return BigDecimal.valueOf(((Number) value).doubleValue()); + default: + throw new IllegalArgumentException("Unsupported type: " + _type); + } + } + assert value == null : "Input value must be a Number or null for numeric endomorphism functions, " + + "but found " + value + " of type " + value.getClass().getName(); + return null; + }) + .collect(Collectors.toList()); + } + + @Override + public Object getResult(NullHandling nullHandling) { + Object baseValue = _base.getResult(nullHandling); + if (baseValue instanceof Number) { + switch (_type) { + case INT: + return ((Number) baseValue).intValue(); + case LONG: + return ((Number) baseValue).longValue(); + case FLOAT: + return ((Number) baseValue).floatValue(); + case DOUBLE: + return ((Number) baseValue).doubleValue(); + case BIG_DECIMAL: + return BigDecimal.valueOf(((Number) baseValue).doubleValue()); + default: + throw new IllegalArgumentException("Unsupported type: " + _type); + } + } + assert baseValue == null : "Input value must be a Number or null for numeric endomorphism functions, " + + "but found " + baseValue + " of type " + baseValue.getClass().getName(); + return null; + } + } + } + + class MultiBuilder implements UdfExampleBuilder { + private final Map> _cases = new HashMap<>(); + + public MultiBuilder and(UdfExampleBuilder testCaseGenerator) { + Map> newCases = testCaseGenerator.generateExamples(); + for (Map.Entry> entry : newCases.entrySet()) { + _cases.merge(entry.getKey(), entry.getValue(), (existing, newer) -> { + existing.addAll(newer); + return existing; + }); + } + return this; + } + + @Override + public Map> generateExamples() { + return _cases; + } + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/udf/UdfParameter.java b/pinot-core/src/main/java/org/apache/pinot/core/udf/UdfParameter.java new file mode 100644 index 000000000000..84bdcaab4920 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/udf/UdfParameter.java @@ -0,0 +1,178 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.core.udf; + +import java.util.ArrayList; +import java.util.List; +import java.util.Objects; +import javax.annotation.Nullable; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; + +/// A parameter on the signature of a [UDF][Udf]. +/// +/// This class contains the DataType of the parameter and whether it is multivalued. +public class UdfParameter { + private final String _name; + private final FieldSpec.DataType _dataType; + private final boolean _multivalued; + private final List _constraints; + @Nullable + private final String _description; + + private UdfParameter( + String name, + FieldSpec.DataType dataType, + boolean multivalued, + List constraints, + String description) { + _name = name; + _dataType = dataType; + _multivalued = multivalued; + _constraints = constraints; + _description = description; + } + + public static UdfParameter result(FieldSpec.DataType dataType) { + return of("result", dataType); + } + + public static UdfParameter of(String name, FieldSpec.DataType dataType) { + return new UdfParameter(name, dataType, false, List.of(), null); + } + + public UdfParameter asMultiValued() { + if (_multivalued) { + return this; + } + return new UdfParameter(_name, _dataType, true, _constraints, _description); + } + + public String getName() { + return _name; + } + + public FieldSpec.DataType getDataType() { + return _dataType; + } + + public boolean isMultivalued() { + return _multivalued; + } + + public boolean isLiteralOnly() { + return _constraints.contains(LiteralConstraint.INSTANCE); + } + + @Nullable + public String getDescription() { + return _description; + } + + public UdfParameter withDescription(String description) { + return new UdfParameter(_name, _dataType, _multivalued, _constraints, description); + } + + public UdfParameter asLiteralOnly() { + if (isLiteralOnly()) { + return this; + } + return new UdfParameter(_name, _dataType, _multivalued, List.of(LiteralConstraint.INSTANCE), _description); + } + + /// Returns the list of constraints that this parameter has. + /// + /// This can be useful to define specific static constraints on the parameter, such as having a specific index. + public List getConstraints() { + return _constraints; + } + + public UdfParameter constrainedWith(Constraint constraint) { + List newConstraints; + if (_constraints.isEmpty()) { + newConstraints = List.of(constraint); + } else { + newConstraints = new ArrayList<>(_constraints.size() + 1); + newConstraints.addAll(_constraints); + newConstraints.add(constraint); + } + return new UdfParameter(_name, _dataType, _multivalued, newConstraints, _description); + } + + @Override + public boolean equals(Object o) { + if (!(o instanceof UdfParameter)) { + return false; + } + UdfParameter that = (UdfParameter) o; + return isMultivalued() == that.isMultivalued() && Objects.equals(_name, that._name) + && getDataType() == that.getDataType() && Objects.equals(getConstraints(), that.getConstraints()) + && Objects.equals(getDescription(), that.getDescription()); + } + + @Override + public int hashCode() { + return Objects.hash(_name, getDataType(), isMultivalued(), getConstraints(), getDescription()); + } + + @Override + public String toString() { + return _name + ": " + getTypeString(); + } + + public String getTypeString() { + String type; + if (isMultivalued()) { + type = "ARRAY(" + getDataType().name().toLowerCase() + ")"; + } else { + type = getDataType().name().toLowerCase(); + } + return type; + } + + /// A constraint on a UDF parameter. + /// + /// This can be used to define specific constraints on the parameter, such as requiring it to be a literal value, + /// or to have a specific index. + /// + /// Remember equals and hashcode should be implemented. Otherwise UdfParameter.equals() may fail unexpectedly. + public interface Constraint { + void updateTableConfig(TableConfigBuilder tableConfigBuilder, String columnName); + } + + public static class LiteralConstraint implements Constraint { + public static final LiteralConstraint INSTANCE = new LiteralConstraint(); + private LiteralConstraint() { + } + + @Override + public void updateTableConfig(TableConfigBuilder tableConfigBuilder, String columnName) { + } + + @Override + public int hashCode() { + return getClass().hashCode(); + } + + @Override + public boolean equals(Object obj) { + return obj instanceof LiteralConstraint; + } + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/udf/UdfSignature.java b/pinot-core/src/main/java/org/apache/pinot/core/udf/UdfSignature.java new file mode 100644 index 000000000000..5536ee2c39f6 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/udf/UdfSignature.java @@ -0,0 +1,73 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.core.udf; + +import java.util.Arrays; +import java.util.List; +import java.util.stream.Collectors; + + +/// The signature of a User Defined Function ([UDF][Udf]) used to define examples. +/// +/// This is a simplified version of the signature defined by Calcite, which is quite more expressive (ie the result +/// type can depend on the parameters). +public abstract class UdfSignature { + public abstract List getParameters(); + + public abstract UdfParameter getReturnType(); + + public int getArity() { + return getParameters().size(); + } + + @Override + public boolean equals(Object o) { + if (!(o instanceof UdfSignature)) { + return false; + } + UdfSignature other = (UdfSignature) o; + return getReturnType().equals(other.getReturnType()) && getParameters().equals(other.getParameters()); + } + + @Override + public int hashCode() { + return 31 * getParameters().hashCode() + getReturnType().hashCode(); + } + + @Override + public String toString() { + return getParameters().stream() + .map(UdfParameter::toString) + .collect(Collectors.joining(", ", "(", ")")) + " -> " + getReturnType().getTypeString(); + } + + public static UdfSignature of(UdfParameter... parametersAndReturnType) { + return new UdfSignature() { + @Override + public List getParameters() { + return Arrays.asList(parametersAndReturnType).subList(0, parametersAndReturnType.length - 1); + } + + @Override + public UdfParameter getReturnType() { + return parametersAndReturnType[parametersAndReturnType.length - 1]; + } + }; + } +} diff --git a/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTest.java b/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTest.java index 21f3dc3c6b7c..1c3e3e68bc6c 100644 --- a/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTest.java +++ b/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTest.java @@ -662,12 +662,26 @@ protected boolean injectTombstones() { return false; } - protected void createAndUploadSegmentFromFile(TableConfig tableConfig, Schema schema, String dataFilePath, + protected void createAndUploadSegmentFromClasspath(TableConfig tableConfig, Schema schema, String dataFilePath, FileFormat fileFormat, long expectedNoOfDocs, long timeoutMs) throws Exception { URL dataPathUrl = getClass().getClassLoader().getResource(dataFilePath); assert dataPathUrl != null; File file = new File(dataPathUrl.getFile()); + createAndUploadSegmentFromFile(tableConfig, schema, file, fileFormat, expectedNoOfDocs, timeoutMs); + } + + /// @deprecated use createAndUploadSegmentFromClasspath instead, given what this class does is to look for + /// dataFilePath on the classpath + @Deprecated + protected void createAndUploadSegmentFromFile(TableConfig tableConfig, Schema schema, String dataFilePath, + FileFormat fileFormat, long expectedNoOfDocs, long timeoutMs) throws Exception { + createAndUploadSegmentFromClasspath(tableConfig, schema, dataFilePath, fileFormat, expectedNoOfDocs, timeoutMs); + } + + protected void createAndUploadSegmentFromFile(TableConfig tableConfig, Schema schema, File file, + FileFormat fileFormat, long expectedNoOfDocs, long timeoutMs) throws Exception { + TestUtils.ensureDirectoriesExistAndEmpty(_segmentDir, _tarDir); ClusterIntegrationTestUtils.buildSegmentFromFile(file, tableConfig, schema, "%", _segmentDir, _tarDir, fileFormat); uploadSegments(tableConfig.getTableName(), _tarDir); diff --git a/pinot-integration-tests/pom.xml b/pinot-integration-tests/pom.xml index 8e76f3ea6d6e..ad09986fa7ac 100644 --- a/pinot-integration-tests/pom.xml +++ b/pinot-integration-tests/pom.xml @@ -222,6 +222,11 @@ test-jar test + + org.apache.pinot + pinot-udf-test + test + org.testng diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/MultiStageEngineIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/MultiStageEngineIntegrationTest.java index 4b0b97b51b74..11c89995f2e5 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/MultiStageEngineIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/MultiStageEngineIntegrationTest.java @@ -1618,7 +1618,7 @@ public void testLookupJoin() throws Exception { TenantConfig tenantConfig = new TenantConfig(getBrokerTenant(), getServerTenant(), null); tableConfig.setTenantConfig(tenantConfig); addTableConfig(tableConfig); - createAndUploadSegmentFromFile(tableConfig, lookupTableSchema, DIM_TABLE_DATA_PATH, FileFormat.CSV, + createAndUploadSegmentFromClasspath(tableConfig, lookupTableSchema, DIM_TABLE_DATA_PATH, FileFormat.CSV, DIM_NUMBER_OF_RECORDS, 60_000); // Compare total rows in the primary table with number of rows in the result of the join with lookup table diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/udf/AvroSink.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/udf/AvroSink.java new file mode 100644 index 000000000000..a97f949edcf3 --- /dev/null +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/udf/AvroSink.java @@ -0,0 +1,191 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.integration.tests.udf; + +import java.io.File; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.util.Arrays; +import java.util.Comparator; +import java.util.EnumMap; +import java.util.List; +import java.util.stream.Collectors; +import org.apache.avro.LogicalTypes; +import org.apache.avro.Schema; +import org.apache.avro.SchemaBuilder; +import org.apache.avro.file.DataFileWriter; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericDatumWriter; +import org.apache.avro.generic.GenericRecord; +import org.apache.pinot.plugin.inputformat.avro.AvroUtils; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.readers.GenericRow; + + +/// Like SegmentProcessorAvroUtils, but: +/// - Assumes generic rows will use normal java types instead of stored types (ie booleans are not stored as ints) +/// - Uses normal avro types instead of stored types (ie booleans are written as booleans, big decimals as bytes, +/// etc.) +/// - Internally manages the Avro objects +/// - Defines a sink like external interface, +public class AvroSink implements AutoCloseable { + + private static final EnumMap NOT_NULL_SCALAR_MAP; + private static final EnumMap NULL_SCALAR_MAP; + private static final EnumMap NOT_NULL_MULTI_VALUE_MAP; + private static final EnumMap NULL_MULTI_VALUE_MAP; + + static { + NOT_NULL_SCALAR_MAP = new EnumMap<>(FieldSpec.DataType.class); + NULL_SCALAR_MAP = new EnumMap<>(FieldSpec.DataType.class); + NOT_NULL_MULTI_VALUE_MAP = new EnumMap<>(FieldSpec.DataType.class); + NULL_MULTI_VALUE_MAP = new EnumMap<>(FieldSpec.DataType.class); + + Schema nullSchema = Schema.create(Schema.Type.NULL); + for (FieldSpec.DataType value : FieldSpec.DataType.values()) { + switch (value) { + case INT: + addType(value, Schema.create(Schema.Type.INT), nullSchema); + break; + case LONG: + addType(value, Schema.create(Schema.Type.LONG), nullSchema); + break; + case FLOAT: + addType(value, Schema.create(Schema.Type.FLOAT), nullSchema); + break; + case DOUBLE: + addType(value, Schema.create(Schema.Type.DOUBLE), nullSchema); + break; + case STRING: + case JSON: + addType(value, Schema.create(Schema.Type.STRING), nullSchema); + break; + case BIG_DECIMAL: + Schema bigDecimal = LogicalTypes.bigDecimal().addToSchema(SchemaBuilder.builder().bytesType()); + addType(value, bigDecimal, nullSchema); + break; + case BYTES: + addType(value, Schema.create(Schema.Type.BYTES), nullSchema); + break; + case BOOLEAN: + addType(value, Schema.create(Schema.Type.BOOLEAN), nullSchema); + break; + case TIMESTAMP: + Schema timestampMillis = LogicalTypes.timestampMillis().addToSchema(SchemaBuilder.builder().longType()); + addType(value, timestampMillis, nullSchema); + break; + case MAP: + case LIST: + case STRUCT: + case UNKNOWN: + // Types we know we don't support in AVRO + break; + default: + throw new RuntimeException("Unsupported data type: " + value); + } + } + } + + private final Schema _avroSchema; + private final DataFileWriter _dataFileWriter; + + public AvroSink(org.apache.pinot.spi.data.Schema schema, File tempFile) + throws IOException { + _avroSchema = convertPinotSchemaToAvroSchema(schema); + + GenericDatumWriter datumWriter = new GenericDatumWriter<>(_avroSchema, AvroUtils.getGenericData()); + _dataFileWriter = new DataFileWriter<>(datumWriter); + _dataFileWriter.create(_avroSchema, tempFile); + } + + public void consume(GenericRow row) + throws IOException { + GenericData.Record record = new GenericData.Record(_avroSchema); + convertGenericRowToAvroRecord(row, record); + _dataFileWriter.append(record); + } + + @Override + public void close() + throws IOException { + _dataFileWriter.close(); + } + + private static void addType( + FieldSpec.DataType dataType, Schema scalarSchema, Schema nullSchema) { + NOT_NULL_SCALAR_MAP.put(dataType, scalarSchema); + NULL_SCALAR_MAP.put(dataType, Schema.createUnion(scalarSchema, nullSchema)); + Schema multiValueSchema = Schema.createArray(scalarSchema); + NOT_NULL_MULTI_VALUE_MAP.put(dataType, multiValueSchema); + NULL_MULTI_VALUE_MAP.put(dataType, Schema.createUnion(multiValueSchema, nullSchema)); + } + + private static void convertGenericRowToAvroRecord(GenericRow input, + GenericData.Record output) { + for (String field : input.getFieldToValueMap().keySet()) { + Object value = input.getValue(field); + if (value instanceof Object[]) { + output.put(field, Arrays.asList((Object[]) value)); + } else { + if (value instanceof byte[]) { + value = ByteBuffer.wrap((byte[]) value); + } + output.put(field, value); + } + } + } + + /** + * Converts a Pinot schema to an Avro schema + */ + private static Schema convertPinotSchemaToAvroSchema(org.apache.pinot.spi.data.Schema pinotSchema) { + SchemaBuilder.FieldAssembler fieldAssembler = SchemaBuilder.record("record").fields(); + + List orderedFieldSpecs = pinotSchema.getAllFieldSpecs().stream() + .sorted(Comparator.comparing(FieldSpec::getName)) + .collect(Collectors.toList()); + for (FieldSpec fieldSpec : orderedFieldSpecs) { + String name = fieldSpec.getName(); + FieldSpec.DataType dataType = fieldSpec.getDataType(); + + Schema fieldType; + if (fieldSpec.isSingleValueField()) { + if (fieldSpec.isNullable()) { + fieldType = NULL_SCALAR_MAP.get(dataType); + } else { + fieldType = NOT_NULL_SCALAR_MAP.get(dataType); + } + } else { + if (fieldSpec.isNullable()) { + fieldType = NULL_MULTI_VALUE_MAP.get(dataType); + } else { + fieldType = NOT_NULL_MULTI_VALUE_MAP.get(dataType); + } + } + if (fieldType == null) { + throw new RuntimeException("Unsupported data type: " + dataType); + } + + fieldAssembler.name(name) + .type(fieldType) + .noDefault(); + } + return fieldAssembler.endRecord(); + } +} diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/udf/IntegrationUdfTestCluster.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/udf/IntegrationUdfTestCluster.java new file mode 100644 index 000000000000..299584cbf194 --- /dev/null +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/udf/IntegrationUdfTestCluster.java @@ -0,0 +1,227 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.integration.tests.udf; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.node.ObjectNode; +import java.io.File; +import java.io.IOException; +import java.io.UncheckedIOException; +import java.util.Iterator; +import java.util.List; +import java.util.Properties; +import java.util.stream.Stream; +import org.apache.pinot.client.ConnectionFactory; +import org.apache.pinot.client.grpc.GrpcConnection; +import org.apache.pinot.client.grpc.GrpcUtils; +import org.apache.pinot.common.proto.Broker; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.controller.helix.ControllerRequestClient; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.integration.tests.BaseClusterIntegrationTest; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.FileFormat; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.spi.exception.QueryErrorCode; +import org.apache.pinot.spi.exception.QueryException; +import org.apache.pinot.udf.test.UdfTestCluster; + +/// A [UdfTestCluster] implementation that starts a Pinot cluster for testing UDFs. +public class IntegrationUdfTestCluster extends BaseClusterIntegrationTest + implements UdfTestCluster { + private boolean _started; + private GrpcConnection _grpcConnection; + + public void start() { + if (_started) { + return; + } + _started = true; + // Start Zookeeper + startZk(); + // Start the Pinot cluster + try { + startController(); + startBroker(); + startServer(); + startMinion(); + + _grpcConnection = ConnectionFactory.fromHostListGrpc(new Properties(), + List.of(getBrokerGrpcEndpoint())); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + @Override + public void close() { + // Notice that this shutdown seems to leak some non-daemon threads. + // This means we need to kill the JVM to stop the test. + // That is fine for TestNG, but not for other callers (like main methods) + if (!_started) { + return; + } + _started = false; + // Stop the Pinot cluster + try { + if (_grpcConnection != null) { + _grpcConnection.close(); + } + + stopMinion(); + stopServer(); + stopBroker(); + stopController(); + } catch (Exception e) { + throw new RuntimeException(e); + } + // Stop Zookeeper + stopZk(); + } + + @Override + public void addTable(Schema schema, TableConfig tableConfig) { + ControllerRequestClient client = getControllerRequestClient(); + try { + client.addSchema(schema); + client.addTableConfig(tableConfig); + } catch (Exception e) { + throw new RuntimeException("Failed to add table: " + tableConfig.getTableName(), e); + } + } + + @Override + public void addRows(String tableName, Schema schema, Stream rows) { + try { + ControllerRequestClient client = getControllerRequestClient(); + TableConfig tableConfig = client.getTableConfig(tableName, TableType.OFFLINE); + + int numRows = 0; + File tempFile = File.createTempFile(tableName, ".avro"); + + try (Stream closeMe = rows; + AvroSink sink = new AvroSink(schema, tempFile)) { + + Iterator it = rows.iterator(); + while (it.hasNext()) { + GenericRow row = it.next(); + try { + sink.consume(row); + } catch (Exception e) { + throw new RuntimeException("Failed to add row " + row + " to table: " + tableName, e); + } + numRows++; + } + } + + try { + createAndUploadSegmentFromFile(tableConfig, schema, tempFile, FileFormat.AVRO, numRows, 10_000); + } catch (Exception e) { + throw new RuntimeException("Failed to add rows to table: " + tableName, e); + } finally { + if (tempFile.exists()) { + tempFile.delete(); + } + } + } catch (IOException e) { + throw new UncheckedIOException(e); + } + } + + @Override + public Iterator query(ExecutionContext context, String sql) { + Iterator response = queryGrpc(context, sql); + // Process metadata + if (!response.hasNext()) { + throw new RuntimeException("No response received from Pinot"); + } + ObjectNode brokerResponseJson = null; + try { + brokerResponseJson = GrpcUtils.extractMetadataJson(response.next()); + } catch (IOException e) { + throw new UncheckedIOException("Cannot extract metadata from GRPC request", e); + } + // Directly return when there is exception + if (brokerResponseJson.has("exceptions") && !brokerResponseJson.get("exceptions").isEmpty()) { + JsonNode exceptions = brokerResponseJson.get("exceptions"); + QueryException toThrow = null; + try { + if (exceptions.isArray() && !exceptions.isEmpty()) { + // If there is only one exception, throw it directly + JsonNode exceptionNode = exceptions.get(0); + QueryErrorCode errorCode = QueryErrorCode.fromErrorCode(exceptionNode.get("errorCode").asInt()); + toThrow = errorCode.asException(exceptionNode.get("message").toString()); + } + } catch (Exception e) { + // nothing to do here + } + throw toThrow != null ? toThrow : new RuntimeException(exceptions.toString()); + } + // Process schema + if (!response.hasNext()) { + throw new RuntimeException("No schema found in the response"); + } + DataSchema dataSchema; + try { + dataSchema = GrpcUtils.extractSchema(response.next()); + } catch (IOException e) { + throw new UncheckedIOException("Cannot extract schema from GRPC request", e); + } + return new Iterator() { + private Iterator _currentBlock = null; + + @Override + public boolean hasNext() { + return response.hasNext() || (_currentBlock != null && _currentBlock.hasNext()); + } + + @Override + public GenericRow next() { + if (_currentBlock != null && _currentBlock.hasNext()) { + Object[] row = _currentBlock.next(); + GenericRow genericRow = new GenericRow(); + for (int i = 0; i < dataSchema.size(); i++) { + String columnName = dataSchema.getColumnName(i); + genericRow.putValue(columnName, row[i]); + } + return genericRow; + } + try { + _currentBlock = GrpcUtils.extractResultTable(response.next(), dataSchema).getRows().iterator(); + } catch (Exception e) { + throw new RuntimeException("Failed to extract row from GRPC response", e); + } + return next(); + } + }; + } + + public Iterator queryGrpc(ExecutionContext context, String sql) { + String prefix = ""; + if (context.getNullHandlingMode() == UdfExample.NullHandling.ENABLED) { + prefix = "SET enableNullHandling=true;\n" + prefix; + } + if (context.isUseMultistageEngine()) { + prefix = "SET useMultistageEngine=true;\n" + prefix; + } + return _grpcConnection.executeWithIterator(prefix + sql); + } +} diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/udf/UdfTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/udf/UdfTest.java new file mode 100644 index 000000000000..2c44a44986fd --- /dev/null +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/udf/UdfTest.java @@ -0,0 +1,444 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.integration.tests.udf; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.databind.MapperFeature; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.json.JsonMapper; +import com.fasterxml.jackson.dataformat.yaml.YAMLFactory; +import java.io.File; +import java.io.FileWriter; +import java.io.IOException; +import java.io.Writer; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.TreeMap; +import java.util.TreeSet; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.stream.Collectors; +import javax.annotation.Nullable; +import org.apache.commons.io.output.StringBuilderWriter; +import org.apache.hadoop.util.Sets; +import org.apache.pinot.common.function.FunctionRegistry; +import org.apache.pinot.common.function.PinotScalarFunction; +import org.apache.pinot.core.operator.transform.function.TransformFunction; +import org.apache.pinot.core.operator.transform.function.TransformFunctionFactory; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.udf.test.UdfReporter; +import org.apache.pinot.udf.test.UdfTestFramework; +import org.apache.pinot.udf.test.UdfTestResult; +import org.assertj.core.api.Assertions; +import org.assertj.core.util.diff.DiffUtils; +import org.assertj.core.util.diff.Patch; +import org.testcontainers.shaded.com.google.common.util.concurrent.MoreExecutors; +import org.testng.Assert; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +/// A TestNG test class that uses snapshot files to verify the results of the UDF test framework. +/// +/// Snapshots are stored in the `src/test/resources/udf-test-results` directory and can be regenerated by running the +/// main method of this class. +/// +/// There are different test that use these snapshots. Read their documentation for more details. +public class UdfTest { + private static final ObjectMapper MAPPER = JsonMapper.builder(new YAMLFactory()) + .configure(MapperFeature.SORT_PROPERTIES_ALPHABETICALLY, true) + .build(); + public static final String ALL_FUNCTIONS_YAML = "all-functions.yaml"; + private IntegrationUdfTestCluster _cluster; + private UdfTestFramework _framework; + private ExecutorService _executorService; + + + @BeforeClass + public void setUp() { + _cluster = new IntegrationUdfTestCluster(); + _cluster.start(); + + _executorService = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors()); + + _framework = UdfTestFramework.fromServiceLoader(_cluster, _executorService); + _framework.startUp(); + } + + @AfterClass + public void tearDown() { + if (_cluster != null) { + _cluster.close(); + } + if (_executorService != null) { + MoreExecutors.shutdownAndAwaitTermination(_executorService, 10, java.util.concurrent.TimeUnit.SECONDS); + } + } + + @DataProvider + public Object[][] udfDataProvider() { + return _framework.getUdfs().stream() + .map(udf -> new Object[]{udf}) + .toArray(Object[][]::new); + } + + /// For each UDF, this test runs the UDF test framework and compares the results with the snapshot file. + /// + /// This test may fail in different situations: + /// 1. A new UDF was added, but the snapshot file was not updated. In this case, the test will fail saying that you + /// need to run the snapshot generation first. + /// 2. The function implementation was changed. In case the semantic change was intentional, you can regenerate + /// the snapshots by running the main method of this class, but you should indicate that the change is breaking + /// change in the pull request and probably also in the commit message. + @Test(dataProvider = "udfDataProvider") + public void testUdf(Udf udf) + throws IOException, InterruptedException { + File snapshotDir = Path.of("src", "test", "resources", "udf-test-results").toFile(); + File snapshotFile = new File(snapshotDir, udf.getMainCanonicalName() + ".yaml"); + if (!snapshotFile.exists()) { + Assert.fail("Snapshot file for UDF " + udf.getMainName() + " does not exist. " + + "Please run the snapshot generation first."); + } + + UdfTestResult.ByScenario actualResult = _framework.execute(udf); + + // Diff failed with an error. Maybe diff CLI is not available. Lets compare the content directly. + List snapshotLines = Files.readAllLines(snapshotFile.toPath()); + List actualLines; + try (StringBuilderWriter writer = new StringBuilderWriter()) { + serializeScenario(writer, actualResult); + actualLines = Arrays.asList(writer.toString().split("\n")); + } + Patch diff = DiffUtils.diff(snapshotLines, actualLines); + + Assertions.assertThat(diff.getDeltas()) + .describedAs("Differences found between the actual and expected UDF test results for %s. " + + "Expected version was %s. " + + "If you think the new solution is correct, please regenerate the examples and add them to the " + + "repository", + udf.getMainName(), snapshotFile.getAbsoluteFile()) + .isEmpty(); + } + + /// This test checks that the snapshot directory does not contain any files that are not used by the UDF test. + /// This could happen if a UDF was removed or renamed, but the snapshot file was not updated. + @Test + public void failWhenSnapshotNotUsed() { + File snapshotDir = Path.of("src", "test", "resources", "udf-test-results").toFile(); + if (!snapshotDir.exists()) { + Assert.fail("Snapshot directory does not exist. Please run the snapshot generation first."); + } + + // Check if the snapshot directory is empty + File[] files = getSnapshotFiles(snapshotDir); + if (files.length == 0) { + return; + } + Set udfNames = _framework.getUdfs().stream() + .map(Udf::getMainCanonicalName) + .collect(Collectors.toSet()); + + List orphanedFiles = new ArrayList<>(); + for (File file : files) { + String udfName = file.getName().replace(".yaml", ""); + if (!udfNames.contains(udfName)) { + orphanedFiles.add(file); + } + } + Assertions.assertThat(orphanedFiles) + .describedAs("The snapshot directory %s contains files that are not used by the UDF test framework. " + + "Please remove them or regenerate the snapshots.", + snapshotDir.getAbsolutePath()) + .isEmpty(); + } + + /// This test checks that the all-functions.yaml file is updated with the latest UDFs. + /// This file keeps track of all the functions that are registered in the system, including scalar and transform + /// functions, and their corresponding UDFs. + /// + /// This could fail if a new function was added to the system, but the all-functions.yaml file was not updated or + /// if a UDF was removed or renamed, but the all-functions.yaml file was not updated. + @Test + public void failWhenAllFunctionsYamlNotUpdated() + throws IOException { + File snapshotDir = Path.of("src", "test", "resources", "udf-test-results").toFile(); + if (!snapshotDir.exists()) { + Assert.fail("Snapshot directory " + snapshotDir.getAbsolutePath() + " does not exist. " + + "Please run the snapshot generation first."); + } + + File allFunctionsFile = new File(snapshotDir, ALL_FUNCTIONS_YAML); + if (!allFunctionsFile.exists()) { + Assert.fail("The all-functions.yaml file " + allFunctionsFile.getAbsolutePath() + " does not exist in the " + + "snapshot directory. Please run the snapshot generation first."); + } + + List expectedAllFunctionsLines = Files.readAllLines(allFunctionsFile.toPath()); + List actualAllFunctionsLines; + try (StringBuilderWriter writer = new StringBuilderWriter()) { + generateAllFunctionsYaml(writer); + actualAllFunctionsLines = Arrays.asList(writer.toString().split("\n")); + } + + Patch diff = DiffUtils.diff(actualAllFunctionsLines, expectedAllFunctionsLines); + Assertions.assertThat(diff.getDeltas()) + .describedAs("Differences found between the actual and expected all-functions.yaml file. " + + "If you think the new solution is correct, please regenerate the examples and add them to the " + + "repository", + allFunctionsFile.getAbsolutePath()) + .isEmpty(); + } + + private File[] getSnapshotFiles(File snapshotDir) { + File[] files = snapshotDir.listFiles(file -> file.isFile() + && file.getName().endsWith(".yaml") + && !file.getName().equals(ALL_FUNCTIONS_YAML)); + if (files == null) { + return new File[0]; + } + return files; + } + + /// Returns a map whose keys are all the function names (including scalar and transform functions) and values are the + /// Udf of that function, or null if the function is not UDF registered. + private TreeMap generateUdfMapForAllFunctions() { + Map scalarFunctions = FunctionRegistry.getFunctions(); + Map> transformFunctions = TransformFunctionFactory.getAllFunctions(); + + TreeMap udfMap = new TreeMap<>(); + for (Udf udf : _framework.getUdfs()) { + for (String name : udf.getAllCanonicalNames()) { + AllFunctionsValue funValue = new AllFunctionsValue(udf.getClass(), transformFunctions.get(name), + scalarFunctions.get(name)); + udfMap.put(name, funValue); + } + } + + Set allFunctions = new TreeSet<>(Sets.union(scalarFunctions.keySet(), transformFunctions.keySet())); + for (String function : allFunctions) { + if (!udfMap.containsKey(function)) { + AllFunctionsValue funValue = new AllFunctionsValue(null, transformFunctions.get(function), + scalarFunctions.get(function)); + udfMap.put(function, funValue); + } + } + return udfMap; + } + + public static class AllFunctionsValue { + @Nullable + private final Class _udf; + @Nullable + private final String _transform; + @Nullable + private final String _scalar; + + public AllFunctionsValue( + @Nullable Class udf, + @Nullable Class transform, + @Nullable PinotScalarFunction scalarFunction + ) { + _udf = udf; + _transform = transform == null ? null : transform.getCanonicalName(); + _scalar = scalarFunction == null ? null : scalarFunction.getScalarFunctionId(); + } + + @JsonCreator + public AllFunctionsValue( + @Nullable @JsonProperty("udf") Class udf, + @Nullable @JsonProperty("transform") String transform, + @Nullable @JsonProperty("scalar") String scalar) { + _udf = udf; + _transform = transform; + _scalar = scalar; + } + + public Class getUdf() { + return _udf; + } + + @Nullable + public String getTransform() { + return _transform; + } + + @Nullable + public String getScalar() { + return _scalar; + } + } + + /// Runs the framework and creates one file for each UDF with the results of the test. + /// + /// The file will be stored in the given outputDir and will be named as the UDF main function name + /// with the .yaml extension. The content of the file will be a YAML representation of the test results, serialized + /// as defined in [UdfTestResult]. + private void generateSnapshots(File snapshotDir) { + try { + UdfTestResult result = _framework.execute(); + Map results = result.getResults(); + + if (!snapshotDir.exists() && !snapshotDir.mkdirs()) { + throw new IOException("Failed to create output directory: " + snapshotDir.getAbsolutePath()); + } + + // Delete all existing YAML files in the snapshot directory + Arrays.stream(getSnapshotFiles(snapshotDir)) + .forEach(File::delete); + + for (Map.Entry entry : results.entrySet()) { + Udf udf = entry.getKey(); + String udfName = udf.getMainCanonicalName(); + File outFile = new File(snapshotDir, udfName + ".yaml"); + + // Serialize only the DTO for this UDF + UdfTestResult.ByScenario byScenario = entry.getValue(); + + try (FileWriter writer = new FileWriter(outFile)) { + serializeScenario(writer, byScenario); + } + + try (FileWriter writer = new FileWriter(new File(snapshotDir, udfName + ".md"))) { + writer.write(getMdLicenseHeader()); + UdfReporter.reportAsMarkdown(udf, byScenario, writer); + } + } + // Write the map to a file named all-functions.yaml + File allFunctionsFile = new File(snapshotDir, ALL_FUNCTIONS_YAML); + try (FileWriter writer = new FileWriter(allFunctionsFile)) { + generateAllFunctionsYaml(writer); + } + } catch (Exception e) { + throw new RuntimeException("Failed to generate UDF snapshots", e); + } + } + + private void generateAllFunctionsYaml(Writer writer) + throws IOException { + TreeMap udfMap = generateUdfMapForAllFunctions(); + + // Write the license header + writer.write(getYamlLicenseHeader()); + writer.write("# This is a map between actual functions and their optional UDFs.\n" + + "# This file has two purposes:\n" + + "# 1. It is used to have a list of UDFs we should create.\n" + + "# 2. It is used to have a test that fails if a new function is added to the system without a UDF.\n" + + "\n"); + // Serialize the map as a DTO + MAPPER.writeValue(writer, udfMap); + } + + private void serializeScenario(Writer writer, UdfTestResult.ByScenario scenario) + throws IOException { + // Write the license header + writer.write(getYamlLicenseHeader()); + MAPPER.writeValue(writer, scenario.asDto()); + } + + private String getYamlLicenseHeader() { + return "#\n" + + "# Licensed to the Apache Software Foundation (ASF) under one\n" + + "# or more contributor license agreements. See the NOTICE file\n" + + "# distributed with this work for additional information\n" + + "# regarding copyright ownership. The ASF licenses this file\n" + + "# to you under the Apache License, Version 2.0 (the\n" + + "# \"License\"); you may not use this file except in compliance\n" + + "# with the License. You may obtain a copy of the License at\n" + + "#\n" + + "# http://www.apache.org/licenses/LICENSE-2.0\n" + + "#\n" + + "# Unless required by applicable law or agreed to in writing,\n" + + "# software distributed under the License is distributed on an\n" + + "# \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY\n" + + "# KIND, either express or implied. See the License for the\n" + + "# specific language governing permissions and limitations\n" + + "# under the License.\n" + + "#\n" + + "\n" + + "# This file is auto-generated by the UDF test framework. Do not edit it manually.\n" + + "# Use the org.apache.pinot.integration.tests.udfUdfTest.generateSnapshots() method to regenerate it.\n" + + "\n"; + } + + private String getMdLicenseHeader() { + return "\n" + + "\n"; + } + + /// Generates the snapshots used to verify results of the UDF test framework. + /// + /// If an argument is provided, it will be used as the output directory for the snapshots. + /// If no argument is provided, it will try to determine the output directory based on the current working directory. + /// If the current working directory is recognized as the root of the Apache Pinot repository or the + /// pinot-integration-tests folder, it will use the `pinot-integration-tests/src/test/resources/udf-test-results` + /// directory (where the test expects snapshots to be stored). + /// If the current working directory is not recognized, it will throw an exception. + public static void main(String[] args) { + File snapshotDir; + if (args.length == 0) { + Path currentDir = Path.of(System.getProperty("user.dir")); + if (currentDir.endsWith(Path.of("pinot", "pinot-integration-tests"))) { + snapshotDir = Path.of("src", "test", "resources", "udf-test-results").toFile(); + } else if (currentDir.endsWith(Path.of("pinot"))) { + snapshotDir = Path.of("pinot-integration-tests", "src", "test", "resources", "udf-test-results").toFile(); + } else { + throw new IllegalStateException("Cannot determine the snapshot directory. " + + "Please provide the path to the snapshot directory as an argument or execute this process either on " + + "the root of the Apache Pinot directory or on the pinot-integration-tests folder."); + } + } else { + snapshotDir = new File(args[0]); + } + + UdfTest test = new UdfTest(); + test.setUp(); + try { + test.generateSnapshots(snapshotDir); + } finally { + test.tearDown(); + } + // TODO: ClusterTestBasePinotFunctionTestCluster uses BaseClusterIntegrationTest, which leak some threads + // non demon on shutdown, which prevents the JVM from exiting. As a workaround, we call System.exit(0) here. + System.exit(0); + } +} diff --git a/pinot-integration-tests/src/test/resources/udf-test-results/README.md b/pinot-integration-tests/src/test/resources/udf-test-results/README.md new file mode 100644 index 000000000000..c0159701d7f5 --- /dev/null +++ b/pinot-integration-tests/src/test/resources/udf-test-results/README.md @@ -0,0 +1,35 @@ + + +# UDF Test Snapshots + +This directory contains the results of UDF tests run against Apache Pinot. + +IMPORTANT: All the files in this directory are generated by the UDF tests and should not be modified manually. +Instead of modifying these files, you should run the main method of the UdfTest test class to regenerate them. + +There are two main types of files in this directory: +1. `all-functions.yaml`: This file contains a map from function names to their UDF classes. +2. `.yaml`: These files contain the results of running the UDF tests for each function. + The content of these files is the YAML serialization of the `UdfTestResult.BySignature.Dto` class. +3. `.md`: These files contain the results of running the UDF tests for each function in Markdown format. + The content of these files is generated from the `UdfTestResult.BySignature` class and they are intended to be + directly used as documentation, copying these files to the `pinot-docs` repository. + This is not yet tested and the generated .md files will be _gitignored_ by default. + diff --git a/pinot-integration-tests/src/test/resources/udf-test-results/abs.yaml b/pinot-integration-tests/src/test/resources/udf-test-results/abs.yaml new file mode 100644 index 000000000000..b499a73bccdb --- /dev/null +++ b/pinot-integration-tests/src/test/resources/udf-test-results/abs.yaml @@ -0,0 +1,870 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# This file is auto-generated by the UDF test framework. Do not edit it manually. +# Use the org.apache.pinot.integration.tests.udfUdfTest.generateSnapshots() method to regenerate it. + +--- +Ingestion time transformer: + '(arg0: big_decimal) -> big_decimal': + entries: + negative value_big_decimal: + actualResult: 3.0 + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 3.0 + null input_big_decimal: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + positive value_big_decimal: + actualResult: 5.0 + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 5.0 + zero_big_decimal: + actualResult: 0.0 + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 0.0 + error: false + errorMessage: null + '(arg0: double) -> double': + entries: + negative value_double: + actualResult: 3.0 + equivalence: "EQUAL" + error: null + expectedResult: 3.0 + null input_double: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + positive value_double: + actualResult: 5.0 + equivalence: "EQUAL" + error: null + expectedResult: 5.0 + zero_double: + actualResult: 0.0 + equivalence: "EQUAL" + error: null + expectedResult: 0.0 + error: false + errorMessage: null + '(arg0: float) -> float': + entries: + negative value_float: + actualResult: 3.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 3.0 + null input_float: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + positive value_float: + actualResult: 5.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 5.0 + zero_float: + actualResult: 0.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 0.0 + error: false + errorMessage: null + '(arg0: int) -> int': + entries: + negative value_int: + actualResult: 3.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 3 + null input_int: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + positive value_int: + actualResult: 5.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 5 + zero_int: + actualResult: 0.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 0 + error: false + errorMessage: null + '(arg0: long) -> long': + entries: + negative value_long: + actualResult: 3.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 3 + null input_long: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + positive value_long: + actualResult: 5.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 5 + zero_long: + actualResult: 0.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 0 + error: false + errorMessage: null +MSE intermediate stage (with null handling): + '(arg0: big_decimal) -> big_decimal': + entries: + negative value_big_decimal: + actualResult: "3.0" + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 3.0 + null input_big_decimal: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + positive value_big_decimal: + actualResult: "5.0" + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 5.0 + zero_big_decimal: + actualResult: "0.0" + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 0.0 + error: false + errorMessage: null + '(arg0: double) -> double': + entries: + negative value_double: + actualResult: 3.0 + equivalence: "EQUAL" + error: null + expectedResult: 3.0 + null input_double: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + positive value_double: + actualResult: 5.0 + equivalence: "EQUAL" + error: null + expectedResult: 5.0 + zero_double: + actualResult: 0.0 + equivalence: "EQUAL" + error: null + expectedResult: 0.0 + error: false + errorMessage: null + '(arg0: float) -> float': + entries: + negative value_float: + actualResult: 3.0 + equivalence: "EQUAL" + error: null + expectedResult: 3.0 + null input_float: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + positive value_float: + actualResult: 5.0 + equivalence: "EQUAL" + error: null + expectedResult: 5.0 + zero_float: + actualResult: 0.0 + equivalence: "EQUAL" + error: null + expectedResult: 0.0 + error: false + errorMessage: null + '(arg0: int) -> int': + entries: + negative value_int: + actualResult: 3 + equivalence: "EQUAL" + error: null + expectedResult: 3 + null input_int: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + positive value_int: + actualResult: 5 + equivalence: "EQUAL" + error: null + expectedResult: 5 + zero_int: + actualResult: 0 + equivalence: "EQUAL" + error: null + expectedResult: 0 + error: false + errorMessage: null + '(arg0: long) -> long': + entries: + negative value_long: + actualResult: 3 + equivalence: "EQUAL" + error: null + expectedResult: 3 + null input_long: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + positive value_long: + actualResult: 5 + equivalence: "EQUAL" + error: null + expectedResult: 5 + zero_long: + actualResult: 0 + equivalence: "EQUAL" + error: null + expectedResult: 0 + error: false + errorMessage: null +MSE intermediate stage (without null handling): + '(arg0: big_decimal) -> big_decimal': + entries: + negative value_big_decimal: + actualResult: "3.0" + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 3.0 + null input_big_decimal: + actualResult: "0.0" + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 0.0 + positive value_big_decimal: + actualResult: "5.0" + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 5.0 + zero_big_decimal: + actualResult: "0.0" + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 0.0 + error: false + errorMessage: null + '(arg0: double) -> double': + entries: + negative value_double: + actualResult: 3.0 + equivalence: "EQUAL" + error: null + expectedResult: 3.0 + null input_double: + actualResult: 0.0 + equivalence: "EQUAL" + error: null + expectedResult: 0.0 + positive value_double: + actualResult: 5.0 + equivalence: "EQUAL" + error: null + expectedResult: 5.0 + zero_double: + actualResult: 0.0 + equivalence: "EQUAL" + error: null + expectedResult: 0.0 + error: false + errorMessage: null + '(arg0: float) -> float': + entries: + negative value_float: + actualResult: 3.0 + equivalence: "EQUAL" + error: null + expectedResult: 3.0 + null input_float: + actualResult: 0.0 + equivalence: "EQUAL" + error: null + expectedResult: 0.0 + positive value_float: + actualResult: 5.0 + equivalence: "EQUAL" + error: null + expectedResult: 5.0 + zero_float: + actualResult: 0.0 + equivalence: "EQUAL" + error: null + expectedResult: 0.0 + error: false + errorMessage: null + '(arg0: int) -> int': + entries: + negative value_int: + actualResult: 3 + equivalence: "EQUAL" + error: null + expectedResult: 3 + null input_int: + actualResult: 0 + equivalence: "EQUAL" + error: null + expectedResult: 0 + positive value_int: + actualResult: 5 + equivalence: "EQUAL" + error: null + expectedResult: 5 + zero_int: + actualResult: 0 + equivalence: "EQUAL" + error: null + expectedResult: 0 + error: false + errorMessage: null + '(arg0: long) -> long': + entries: + negative value_long: + actualResult: 3 + equivalence: "EQUAL" + error: null + expectedResult: 3 + null input_long: + actualResult: 0 + equivalence: "EQUAL" + error: null + expectedResult: 0 + positive value_long: + actualResult: 5 + equivalence: "EQUAL" + error: null + expectedResult: 5 + zero_long: + actualResult: 0 + equivalence: "EQUAL" + error: null + expectedResult: 0 + error: false + errorMessage: null +SSE predicate (with null handling): + '(arg0: big_decimal) -> big_decimal': + entries: + negative value_big_decimal: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null input_big_decimal: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + positive value_big_decimal: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + zero_big_decimal: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null + '(arg0: double) -> double': + entries: + negative value_double: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null input_double: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + positive value_double: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + zero_double: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null + '(arg0: float) -> float': + entries: + negative value_float: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null input_float: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + positive value_float: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + zero_float: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null + '(arg0: int) -> int': + entries: + negative value_int: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null input_int: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + positive value_int: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + zero_int: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null + '(arg0: long) -> long': + entries: + negative value_long: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null input_long: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + positive value_long: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + zero_long: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE predicate (without null handling): + '(arg0: big_decimal) -> big_decimal': + entries: + negative value_big_decimal: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null input_big_decimal: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + positive value_big_decimal: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + zero_big_decimal: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null + '(arg0: double) -> double': + entries: + negative value_double: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null input_double: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + positive value_double: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + zero_double: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null + '(arg0: float) -> float': + entries: + negative value_float: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null input_float: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + positive value_float: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + zero_float: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null + '(arg0: int) -> int': + entries: + negative value_int: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null input_int: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + positive value_int: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + zero_int: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null + '(arg0: long) -> long': + entries: + negative value_long: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null input_long: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + positive value_long: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + zero_long: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE projection (with null handling): + '(arg0: big_decimal) -> big_decimal': + entries: + negative value_big_decimal: + actualResult: "3" + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 3.0 + null input_big_decimal: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + positive value_big_decimal: + actualResult: "5" + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 5.0 + zero_big_decimal: + actualResult: "0" + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 0.0 + error: false + errorMessage: null + '(arg0: double) -> double': + entries: + negative value_double: + actualResult: 3.0 + equivalence: "EQUAL" + error: null + expectedResult: 3.0 + null input_double: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + positive value_double: + actualResult: 5.0 + equivalence: "EQUAL" + error: null + expectedResult: 5.0 + zero_double: + actualResult: 0.0 + equivalence: "EQUAL" + error: null + expectedResult: 0.0 + error: false + errorMessage: null + '(arg0: float) -> float': + entries: + negative value_float: + actualResult: 3.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 3.0 + null input_float: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + positive value_float: + actualResult: 5.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 5.0 + zero_float: + actualResult: 0.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 0.0 + error: false + errorMessage: null + '(arg0: int) -> int': + entries: + negative value_int: + actualResult: 3.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 3 + null input_int: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + positive value_int: + actualResult: 5.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 5 + zero_int: + actualResult: 0.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 0 + error: false + errorMessage: null + '(arg0: long) -> long': + entries: + negative value_long: + actualResult: 3.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 3 + null input_long: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + positive value_long: + actualResult: 5.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 5 + zero_long: + actualResult: 0.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 0 + error: false + errorMessage: null +SSE projection (without null handling): + '(arg0: big_decimal) -> big_decimal': + entries: + negative value_big_decimal: + actualResult: "3" + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 3.0 + null input_big_decimal: + actualResult: "0" + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 0.0 + positive value_big_decimal: + actualResult: "5" + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 5.0 + zero_big_decimal: + actualResult: "0" + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 0.0 + error: false + errorMessage: null + '(arg0: double) -> double': + entries: + negative value_double: + actualResult: 3.0 + equivalence: "EQUAL" + error: null + expectedResult: 3.0 + null input_double: + actualResult: 0.0 + equivalence: "EQUAL" + error: null + expectedResult: 0.0 + positive value_double: + actualResult: 5.0 + equivalence: "EQUAL" + error: null + expectedResult: 5.0 + zero_double: + actualResult: 0.0 + equivalence: "EQUAL" + error: null + expectedResult: 0.0 + error: false + errorMessage: null + '(arg0: float) -> float': + entries: + negative value_float: + actualResult: 3.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 3.0 + null input_float: + actualResult: 0.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 0.0 + positive value_float: + actualResult: 5.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 5.0 + zero_float: + actualResult: 0.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 0.0 + error: false + errorMessage: null + '(arg0: int) -> int': + entries: + negative value_int: + actualResult: 3.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 3 + null input_int: + actualResult: 0.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 0 + positive value_int: + actualResult: 5.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 5 + zero_int: + actualResult: 0.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 0 + error: false + errorMessage: null + '(arg0: long) -> long': + entries: + negative value_long: + actualResult: 3.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 3 + null input_long: + actualResult: 0.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 0 + positive value_long: + actualResult: 5.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 5 + zero_long: + actualResult: 0.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 0 + error: false + errorMessage: null diff --git a/pinot-integration-tests/src/test/resources/udf-test-results/acos.yaml b/pinot-integration-tests/src/test/resources/udf-test-results/acos.yaml new file mode 100644 index 000000000000..e891b463e9fe --- /dev/null +++ b/pinot-integration-tests/src/test/resources/udf-test-results/acos.yaml @@ -0,0 +1,268 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# This file is auto-generated by the UDF test framework. Do not edit it manually. +# Use the org.apache.pinot.integration.tests.udfUdfTest.generateSnapshots() method to regenerate it. + +--- +Ingestion time transformer: + '(d: double) -> double': + entries: + acos(-1): + actualResult: 3.141592653589793 + equivalence: "EQUAL" + error: null + expectedResult: 3.141592653589793 + acos(0): + actualResult: 1.5707963267948966 + equivalence: "EQUAL" + error: null + expectedResult: 1.5707963267948966 + acos(1): + actualResult: 0.0 + equivalence: "EQUAL" + error: null + expectedResult: 0.0 + larger than 1: + actualResult: NaN + equivalence: "EQUAL" + error: null + expectedResult: NaN + null input: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + smaller than -1: + actualResult: NaN + equivalence: "EQUAL" + error: null + expectedResult: NaN + error: false + errorMessage: null +MSE intermediate stage (with null handling): + '(d: double) -> double': + entries: + acos(-1): + actualResult: 3.141592653589793 + equivalence: "EQUAL" + error: null + expectedResult: 3.141592653589793 + acos(0): + actualResult: 1.5707963267948966 + equivalence: "EQUAL" + error: null + expectedResult: 1.5707963267948966 + acos(1): + actualResult: 0.0 + equivalence: "EQUAL" + error: null + expectedResult: 0.0 + larger than 1: + actualResult: NaN + equivalence: "EQUAL" + error: null + expectedResult: NaN + null input: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + smaller than -1: + actualResult: NaN + equivalence: "EQUAL" + error: null + expectedResult: NaN + error: false + errorMessage: null +MSE intermediate stage (without null handling): + '(d: double) -> double': + entries: + acos(-1): + actualResult: 3.141592653589793 + equivalence: "EQUAL" + error: null + expectedResult: 3.141592653589793 + acos(0): + actualResult: 1.5707963267948966 + equivalence: "EQUAL" + error: null + expectedResult: 1.5707963267948966 + acos(1): + actualResult: 0.0 + equivalence: "EQUAL" + error: null + expectedResult: 0.0 + larger than 1: + actualResult: NaN + equivalence: "EQUAL" + error: null + expectedResult: NaN + null input: + actualResult: 1.5707963267948966 + equivalence: "EQUAL" + error: null + expectedResult: 1.5707963267948966 + smaller than -1: + actualResult: NaN + equivalence: "EQUAL" + error: null + expectedResult: NaN + error: false + errorMessage: null +SSE predicate (with null handling): + '(d: double) -> double': + entries: + acos(-1): + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + acos(0): + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + acos(1): + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + larger than 1: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null input: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + smaller than -1: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE predicate (without null handling): + '(d: double) -> double': + entries: + acos(-1): + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + acos(0): + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + acos(1): + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + larger than 1: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null input: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + smaller than -1: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE projection (with null handling): + '(d: double) -> double': + entries: + acos(-1): + actualResult: 3.141592653589793 + equivalence: "EQUAL" + error: null + expectedResult: 3.141592653589793 + acos(0): + actualResult: 1.5707963267948966 + equivalence: "EQUAL" + error: null + expectedResult: 1.5707963267948966 + acos(1): + actualResult: 0.0 + equivalence: "EQUAL" + error: null + expectedResult: 0.0 + larger than 1: + actualResult: NaN + equivalence: "EQUAL" + error: null + expectedResult: NaN + null input: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + smaller than -1: + actualResult: NaN + equivalence: "EQUAL" + error: null + expectedResult: NaN + error: false + errorMessage: null +SSE projection (without null handling): + '(d: double) -> double': + entries: + acos(-1): + actualResult: 3.141592653589793 + equivalence: "EQUAL" + error: null + expectedResult: 3.141592653589793 + acos(0): + actualResult: 1.5707963267948966 + equivalence: "EQUAL" + error: null + expectedResult: 1.5707963267948966 + acos(1): + actualResult: 0.0 + equivalence: "EQUAL" + error: null + expectedResult: 0.0 + larger than 1: + actualResult: NaN + equivalence: "EQUAL" + error: null + expectedResult: NaN + null input: + actualResult: 1.5707963267948966 + equivalence: "EQUAL" + error: null + expectedResult: 1.5707963267948966 + smaller than -1: + actualResult: NaN + equivalence: "EQUAL" + error: null + expectedResult: NaN + error: false + errorMessage: null diff --git a/pinot-integration-tests/src/test/resources/udf-test-results/adler32.yaml b/pinot-integration-tests/src/test/resources/udf-test-results/adler32.yaml new file mode 100644 index 000000000000..2a65debf2a43 --- /dev/null +++ b/pinot-integration-tests/src/test/resources/udf-test-results/adler32.yaml @@ -0,0 +1,198 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# This file is auto-generated by the UDF test framework. Do not edit it manually. +# Use the org.apache.pinot.integration.tests.udfUdfTest.generateSnapshots() method to regenerate it. + +--- +Ingestion time transformer: + '(input: bytes) -> int': + entries: + empty: + actualResult: 1 + equivalence: "EQUAL" + error: null + expectedResult: 1 + multiple bytes: + actualResult: 1572875 + equivalence: "EQUAL" + error: null + expectedResult: 1572875 + null input: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + single byte: + actualResult: 131074 + equivalence: "EQUAL" + error: null + expectedResult: 131074 + error: false + errorMessage: null +MSE intermediate stage (with null handling): + '(input: bytes) -> int': + entries: + empty: + actualResult: 1 + equivalence: "EQUAL" + error: null + expectedResult: 1 + multiple bytes: + actualResult: 1572875 + equivalence: "EQUAL" + error: null + expectedResult: 1572875 + null input: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + single byte: + actualResult: 131074 + equivalence: "EQUAL" + error: null + expectedResult: 131074 + error: false + errorMessage: null +MSE intermediate stage (without null handling): + '(input: bytes) -> int': + entries: + empty: + actualResult: 1 + equivalence: "EQUAL" + error: null + expectedResult: 1 + multiple bytes: + actualResult: 1572875 + equivalence: "EQUAL" + error: null + expectedResult: 1572875 + null input: + actualResult: 1 + equivalence: "EQUAL" + error: null + expectedResult: 1 + single byte: + actualResult: 131074 + equivalence: "EQUAL" + error: null + expectedResult: 131074 + error: false + errorMessage: null +SSE predicate (with null handling): + '(input: bytes) -> int': + entries: + empty: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + multiple bytes: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null input: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + single byte: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE predicate (without null handling): + '(input: bytes) -> int': + entries: + empty: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + multiple bytes: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null input: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + single byte: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE projection (with null handling): + '(input: bytes) -> int': + entries: + empty: + actualResult: 1 + equivalence: "EQUAL" + error: null + expectedResult: 1 + multiple bytes: + actualResult: 1572875 + equivalence: "EQUAL" + error: null + expectedResult: 1572875 + null input: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + single byte: + actualResult: 131074 + equivalence: "EQUAL" + error: null + expectedResult: 131074 + error: false + errorMessage: null +SSE projection (without null handling): + '(input: bytes) -> int': + entries: + empty: + actualResult: 1 + equivalence: "EQUAL" + error: null + expectedResult: 1 + multiple bytes: + actualResult: 1572875 + equivalence: "EQUAL" + error: null + expectedResult: 1572875 + null input: + actualResult: 1 + equivalence: "EQUAL" + error: null + expectedResult: 1 + single byte: + actualResult: 131074 + equivalence: "EQUAL" + error: null + expectedResult: 131074 + error: false + errorMessage: null diff --git a/pinot-integration-tests/src/test/resources/udf-test-results/all-functions.yaml b/pinot-integration-tests/src/test/resources/udf-test-results/all-functions.yaml new file mode 100644 index 000000000000..b3e68ddf9610 --- /dev/null +++ b/pinot-integration-tests/src/test/resources/udf-test-results/all-functions.yaml @@ -0,0 +1,1952 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# This file is auto-generated by the UDF test framework. Do not edit it manually. +# Use the org.apache.pinot.integration.tests.udfUdfTest.generateSnapshots() method to regenerate it. + +# This is a map between actual functions and their optional UDFs. +# This file has two purposes: +# 1. It is used to have a list of UDFs we should create. +# 2. It is used to have a test that fails if a new function is added to the system without a UDF. + +--- +abs: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.abs}" + transform: "org.apache.pinot.core.operator.transform.function.SingleParamMathTransformFunction.AbsTransformFunction" + udf: "org.apache.pinot.query.runtime.function.AbsUdf" +acos: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.TrigonometricFunctions.acos}" + transform: "org.apache.pinot.core.operator.transform.function.TrigonometricTransformFunctions.AcosTransformFunction" + udf: "org.apache.pinot.query.runtime.function.AcosUdf" +add: + scalar: "org.apache.pinot.common.function.scalar.arithmetic.PlusScalarFunction" + transform: "org.apache.pinot.core.operator.transform.function.AdditionTransformFunction" + udf: "org.apache.pinot.query.runtime.function.PlusUdf" +adler32: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.HashFunctions.adler32}" + transform: null + udf: "org.apache.pinot.query.runtime.function.Adler32Udf" +ago: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.ago}" + transform: null + udf: null +agomv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.agoMV}" + transform: null + udf: null +and: + scalar: null + transform: "org.apache.pinot.core.operator.transform.function.AndOperatorTransformFunction" + udf: "org.apache.pinot.query.runtime.function.AndUdf" +array: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayValueConstructor}" + transform: null + udf: "org.apache.pinot.query.runtime.function.ArrayUdf" +arrayaverage: + scalar: null + transform: "org.apache.pinot.core.operator.transform.function.ArrayAverageTransformFunction" + udf: "org.apache.pinot.query.runtime.function.ArrayAverageUdf" +arrayconcatdouble: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayConcatDouble}" + transform: null + udf: "org.apache.pinot.query.runtime.function.ArrayConcatDoubleUdf" +arrayconcatfloat: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayConcatFloat}" + transform: null + udf: null +arrayconcatint: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayConcatInt}" + transform: null + udf: null +arrayconcatlong: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayConcatLong}" + transform: null + udf: null +arrayconcatstring: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayConcatString}" + transform: null + udf: null +arraycontainsint: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayContainsInt}" + transform: null + udf: null +arraycontainsstring: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayContainsString}" + transform: null + udf: null +arraydistinctint: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayDistinctInt}" + transform: null + udf: null +arraydistinctstring: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayDistinctString}" + transform: null + udf: null +arrayelementatdouble: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayElementAtDouble}" + transform: null + udf: null +arrayelementatfloat: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayElementAtFloat}" + transform: null + udf: null +arrayelementatint: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayElementAtInt}" + transform: null + udf: "org.apache.pinot.query.runtime.function.ArrayElementAtIntUdf" +arrayelementatlong: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayElementAtLong}" + transform: null + udf: null +arrayelementatstring: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayElementAtString}" + transform: null + udf: null +arrayindexesofdouble: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayIndexesOfDouble}" + transform: null + udf: null +arrayindexesoffloat: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayIndexesOfFloat}" + transform: null + udf: null +arrayindexesofint: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayIndexesOfInt}" + transform: null + udf: null +arrayindexesoflong: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayIndexesOfLong}" + transform: null + udf: null +arrayindexesofstring: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayIndexesOfString}" + transform: null + udf: null +arrayindexofint: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayIndexOfInt}" + transform: null + udf: null +arrayindexofstring: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayIndexOfString}" + transform: null + udf: null +arraylength: + scalar: "org.apache.pinot.common.function.scalar.array.ArrayLengthScalarFunction" + transform: "org.apache.pinot.core.operator.transform.function.ArrayLengthTransformFunction" + udf: null +arraymax: + scalar: null + transform: "org.apache.pinot.core.operator.transform.function.ArrayMaxTransformFunction" + udf: "org.apache.pinot.query.runtime.function.ArrayMaxUdf" +arraymin: + scalar: null + transform: "org.apache.pinot.core.operator.transform.function.ArrayMinTransformFunction" + udf: null +arrayremoveint: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayRemoveInt}" + transform: null + udf: null +arrayremovestring: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayRemoveString}" + transform: null + udf: null +arrayreverseint: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayReverseInt}" + transform: null + udf: null +arrayreversestring: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayReverseString}" + transform: null + udf: null +arraysliceint: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arraySliceInt}" + transform: null + udf: null +arrayslicestring: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arraySliceString}" + transform: null + udf: null +arraysortint: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arraySortInt}" + transform: null + udf: null +arraysortstring: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arraySortString}" + transform: null + udf: null +arraysum: + scalar: null + transform: "org.apache.pinot.core.operator.transform.function.ArraySumTransformFunction" + udf: null +arraysumint: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arraySumInt}" + transform: null + udf: null +arraysumlong: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arraySumLong}" + transform: null + udf: null +arraytostring: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.ArrayFunctions.arrayToString,\ + \ 3: org.apache.pinot.common.function.scalar.ArrayFunctions.arrayToString]}" + transform: null + udf: null +arrayunionint: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayUnionInt}" + transform: null + udf: null +arrayunionstring: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayUnionString}" + transform: null + udf: null +arrayvalueconstructor: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.arrayValueConstructor}" + transform: null + udf: "org.apache.pinot.query.runtime.function.ArrayUdf" +asin: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.TrigonometricFunctions.asin}" + transform: "org.apache.pinot.core.operator.transform.function.TrigonometricTransformFunctions.AsinTransformFunction" + udf: null +atan: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.TrigonometricFunctions.atan}" + transform: "org.apache.pinot.core.operator.transform.function.TrigonometricTransformFunctions.AtanTransformFunction" + udf: null +atan2: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.TrigonometricFunctions.atan2}" + transform: "org.apache.pinot.core.operator.transform.function.TrigonometricTransformFunctions.Atan2TransformFunction" + udf: null +avgreduce: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.query.reduce.function.InternalReduceFunctions.avgReduce}" + transform: null + udf: null +base64decode: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DataTypeConversionFunctions.base64Decode}" + transform: null + udf: null +base64encode: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DataTypeConversionFunctions.base64Encode}" + transform: null + udf: null +between: + scalar: "org.apache.pinot.common.function.scalar.comparison.BetweenScalarFunction" + transform: null + udf: null +bigdecimaltobytes: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DataTypeConversionFunctions.bigDecimalToBytes}" + transform: null + udf: null +brokerid: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.InternalFunctions.brokerId}" + transform: null + udf: null +bytestobigdecimal: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DataTypeConversionFunctions.bytesToBigDecimal}" + transform: null + udf: null +bytestohex: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DataTypeConversionFunctions.bytesToHex}" + transform: null + udf: null +byteswapint: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.byteswapInt}" + transform: null + udf: null +byteswaplong: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.byteswapLong}" + transform: null + udf: null +cardinality: + scalar: "org.apache.pinot.common.function.scalar.array.ArrayLengthScalarFunction" + transform: "org.apache.pinot.core.operator.transform.function.ArrayLengthTransformFunction" + udf: null +case: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ObjectFunctions.caseWhen}" + transform: "org.apache.pinot.core.operator.transform.function.CaseTransformFunction" + udf: null +casewhen: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ObjectFunctions.caseWhen}" + transform: null + udf: null +cast: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DataTypeConversionFunctions.cast}" + transform: "org.apache.pinot.core.operator.transform.function.CastTransformFunction" + udf: null +ceil: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.ceil}" + transform: "org.apache.pinot.core.operator.transform.function.SingleParamMathTransformFunction.CeilTransformFunction" + udf: null +ceiling: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.ceil}" + transform: "org.apache.pinot.core.operator.transform.function.SingleParamMathTransformFunction.CeilTransformFunction" + udf: null +chr: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.chr}" + transform: null + udf: null +cid: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.InternalFunctions.cid}" + transform: null + udf: null +cityhash128: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.HashFunctions.cityHash128}" + transform: null + udf: null +cityhash32: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.HashFunctions.cityHash32}" + transform: null + udf: null +cityhash64: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.HashFunctions.cityHash64,\ + \ 2: org.apache.pinot.common.function.scalar.HashFunctions.cityHash64, 3: org.apache.pinot.common.function.scalar.HashFunctions.cityHash64]}" + transform: null + udf: null +clpdecode: + scalar: null + transform: "org.apache.pinot.core.operator.transform.function.CLPDecodeTransformFunction" + udf: null +clpencodedvarsmatch: + scalar: null + transform: "org.apache.pinot.core.operator.transform.function.ClpEncodedVarsMatchTransformFunction" + udf: null +coalesce: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ObjectFunctions.coalesce}" + transform: "org.apache.pinot.core.operator.transform.function.CoalesceTransformFunction" + udf: null +codepoint: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.codepoint}" + transform: null + udf: null +concat: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.StringFunctions.concat,\ + \ 3: org.apache.pinot.common.function.scalar.string.StringFunctions.concat]}" + transform: null + udf: null +concatws: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.string.StringFunctions.concatWS}" + transform: null + udf: null +contains: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.contains}" + transform: null + udf: null +cos: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.TrigonometricFunctions.cos}" + transform: "org.apache.pinot.core.operator.transform.function.TrigonometricTransformFunctions.CosTransformFunction" + udf: null +cosh: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.TrigonometricFunctions.cosh}" + transform: "org.apache.pinot.core.operator.transform.function.TrigonometricTransformFunctions.CoshTransformFunction" + udf: null +cosinedistance: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.VectorFunctions.cosineDistance,\ + \ 3: org.apache.pinot.common.function.scalar.VectorFunctions.cosineDistance]}" + transform: "org.apache.pinot.core.operator.transform.function.VectorTransformFunctions.CosineDistanceTransformFunction" + udf: null +cot: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.TrigonometricFunctions.cot}" + transform: "org.apache.pinot.core.operator.transform.function.TrigonometricTransformFunctions.CotTransformFunction" + udf: null +cpcsketchtostring: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.function.scalar.SketchFunctions.cpcSketchToString}" + transform: null + udf: null +cpcsketchunion: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.core.function.scalar.SketchFunctions.cpcSketchUnion,\ + \ 3: org.apache.pinot.core.function.scalar.SketchFunctions.cpcSketchUnion, 4:\ + \ org.apache.pinot.core.function.scalar.SketchFunctions.cpcSketchUnion, 5: org.apache.pinot.core.function.scalar.SketchFunctions.cpcSketchUnion]}" + transform: null + udf: null +crc32: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.HashFunctions.crc32}" + transform: null + udf: null +crc32c: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.HashFunctions.crc32c}" + transform: null + udf: null +cutfragment: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.cutFragment}" + transform: null + udf: null +cutquerystring: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.cutQueryString}" + transform: null + udf: null +cutquerystringandfragment: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.cutQueryStringAndFragment}" + transform: null + udf: null +cuttofirstsignificantsubdomain: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.cutToFirstSignificantSubdomain}" + transform: null + udf: null +cuttofirstsignificantsubdomainwithwww: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.cutToFirstSignificantSubdomainWithWWW}" + transform: null + udf: null +cuturlparameter: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.cutURLParameter}" + transform: null + udf: null +cuturlparameters: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.cutURLParameters}" + transform: null + udf: null +cutwww: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.cutWWW}" + transform: null + udf: null +dateadd: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.timestampAdd}" + transform: null + udf: null +dateaddmv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.timestampAddMV}" + transform: null + udf: null +datebin: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.dateBin}" + transform: null + udf: null +datediff: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.timestampDiff}" + transform: null + udf: null +datediffmv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.timestampDiffMV}" + transform: null + udf: null +datediffmvreverse: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.timestampDiffMVReverse}" + transform: null + udf: null +datetimeconvert: + scalar: "ArgumentCountBasedScalarFunction{[4: org.apache.pinot.common.function.scalar.DateTimeConvert.dateTimeConvert,\ + \ 5: org.apache.pinot.common.function.scalar.DateTimeConvert.dateTimeConvert]}" + transform: "org.apache.pinot.core.operator.transform.function.DateTimeConversionTransformFunction" + udf: null +datetimeconvertwindowhop: + scalar: null + transform: "org.apache.pinot.core.operator.transform.function.DateTimeConversionHopTransformFunction" + udf: null +datetrunc: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.DateTimeFunctions.dateTrunc,\ + \ 3: org.apache.pinot.common.function.scalar.DateTimeFunctions.dateTrunc, 4: org.apache.pinot.common.function.scalar.DateTimeFunctions.dateTrunc,\ + \ 5: org.apache.pinot.common.function.scalar.DateTimeFunctions.dateTrunc]}" + transform: "org.apache.pinot.core.operator.transform.function.DateTruncTransformFunction" + udf: null +datetruncmv: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.DateTimeFunctions.dateTruncMV,\ + \ 3: org.apache.pinot.common.function.scalar.DateTimeFunctions.dateTruncMV, 4:\ + \ org.apache.pinot.common.function.scalar.DateTimeFunctions.dateTruncMV, 5: org.apache.pinot.common.function.scalar.DateTimeFunctions.dateTruncMV]}" + transform: null + udf: null +day: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.dayOfMonth,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.dayOfMonth]}" + transform: "org.apache.pinot.core.operator.transform.function.DateTimeTransformFunction.DayOfMonth" + udf: null +daymv: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.dayOfMonthMV,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.dayOfMonthMV]}" + transform: null + udf: null +dayofmonth: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.dayOfMonth,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.dayOfMonth]}" + transform: "org.apache.pinot.core.operator.transform.function.DateTimeTransformFunction.DayOfMonth" + udf: null +dayofmonthmv: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.dayOfMonthMV,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.dayOfMonthMV]}" + transform: null + udf: null +dayofweek: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.dayOfWeek,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.dayOfWeek]}" + transform: "org.apache.pinot.core.operator.transform.function.DateTimeTransformFunction.DayOfWeek" + udf: null +dayofweekmv: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.dayOfWeekMV,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.dayOfWeekMV]}" + transform: null + udf: null +dayofyear: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.dayOfYear,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.dayOfYear]}" + transform: "org.apache.pinot.core.operator.transform.function.DateTimeTransformFunction.DayOfYear" + udf: null +dayofyearmv: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.dayOfYear,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.dayOfYear]}" + transform: null + udf: null +decodegeohash: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.GeohashFunctions.decodeGeoHash}" + transform: null + udf: null +decodegeohashlat: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.GeohashFunctions.decodeGeoHashLatitude}" + transform: null + udf: null +decodegeohashlatitude: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.GeohashFunctions.decodeGeoHashLatitude}" + transform: null + udf: null +decodegeohashlon: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.GeohashFunctions.decodeGeoHashLongitude}" + transform: null + udf: null +decodegeohashlongitude: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.GeohashFunctions.decodeGeoHashLongitude}" + transform: null + udf: null +decodeurl: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.decodeUrl}" + transform: null + udf: null +degrees: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.TrigonometricFunctions.degrees}" + transform: "org.apache.pinot.core.operator.transform.function.TrigonometricTransformFunctions.DegreesTransformFunction" + udf: "org.apache.pinot.query.runtime.function.DegreesUdf" +div: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.ArithmeticFunctions.divide,\ + \ 3: org.apache.pinot.common.function.scalar.ArithmeticFunctions.divide]}" + transform: "org.apache.pinot.core.operator.transform.function.DivisionTransformFunction" + udf: null +divide: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.ArithmeticFunctions.divide,\ + \ 3: org.apache.pinot.common.function.scalar.ArithmeticFunctions.divide]}" + transform: "org.apache.pinot.core.operator.transform.function.DivisionTransformFunction" + udf: null +dotproduct: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.VectorFunctions.dotProduct}" + transform: null + udf: null +dow: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.dayOfWeek,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.dayOfWeek]}" + transform: "org.apache.pinot.core.operator.transform.function.DateTimeTransformFunction.DayOfWeek" + udf: null +dowmv: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.dayOfWeekMV,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.dayOfWeekMV]}" + transform: null + udf: null +doy: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.dayOfYear,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.dayOfYear]}" + transform: "org.apache.pinot.core.operator.transform.function.DateTimeTransformFunction.DayOfYear" + udf: null +doymv: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.dayOfYear,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.dayOfYear]}" + transform: null + udf: null +encodegeohash: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.GeohashFunctions.encodeGeoHash}" + transform: null + udf: null +encodeurl: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.encodeUrl}" + transform: null + udf: null +endswith: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.endsWith}" + transform: null + udf: null +endtime: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.InternalFunctions.endTime}" + transform: null + udf: null +equals: + scalar: "org.apache.pinot.common.function.scalar.comparison.EqualsScalarFunction" + transform: "org.apache.pinot.core.operator.transform.function.EqualsTransformFunction" + udf: null +euclideandistance: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.VectorFunctions.euclideanDistance}" + transform: null + udf: null +exp: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.exp}" + transform: "org.apache.pinot.core.operator.transform.function.SingleParamMathTransformFunction.ExpTransformFunction" + udf: null +extract: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.extract}" + transform: "org.apache.pinot.core.operator.transform.function.ExtractTransformFunction" + udf: null +extracturlparameter: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.extractURLParameter}" + transform: null + udf: null +extracturlparameternames: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.extractURLParameterNames}" + transform: null + udf: null +extracturlparameters: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.extractURLParameters}" + transform: null + udf: null +floor: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.floor}" + transform: "org.apache.pinot.core.operator.transform.function.SingleParamMathTransformFunction.FloorTransformFunction" + udf: null +fromascii: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.fromAscii}" + transform: null + udf: null +frombase64: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.fromBase64}" + transform: null + udf: null +frombytes: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.fromBytes}" + transform: null + udf: null +fromdatetime: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.DateTimeFunctions.fromDateTime,\ + \ 3: org.apache.pinot.common.function.scalar.DateTimeFunctions.fromDateTime, 4:\ + \ org.apache.pinot.common.function.scalar.DateTimeFunctions.fromDateTime]}" + transform: null + udf: null +fromdatetimemv: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.DateTimeFunctions.fromDateTimeMV,\ + \ 3: org.apache.pinot.common.function.scalar.DateTimeFunctions.fromDateTimeMV]}" + transform: null + udf: null +fromepochdays: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.fromEpochDays}" + transform: null + udf: null +fromepochdaysbucket: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.fromEpochDaysBucket}" + transform: null + udf: null +fromepochdaysbucketmv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.fromEpochDaysBucketMV}" + transform: null + udf: null +fromepochdaysmv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.fromEpochDaysMV}" + transform: null + udf: null +fromepochhours: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.fromEpochHours}" + transform: null + udf: null +fromepochhoursbucket: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.fromEpochHoursBucket}" + transform: null + udf: null +fromepochhoursbucketmv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.fromEpochHoursBucketMV}" + transform: null + udf: null +fromepochhoursmv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.fromEpochHoursMV}" + transform: null + udf: null +fromepochminutes: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.fromEpochMinutes}" + transform: null + udf: null +fromepochminutesbucket: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.fromEpochMinutesBucket}" + transform: null + udf: null +fromepochminutesbucketmv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.fromEpochMinutesBucketMV}" + transform: null + udf: null +fromepochminutesmv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.fromEpochMinutesMV}" + transform: null + udf: null +fromepochseconds: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.fromEpochSeconds}" + transform: null + udf: null +fromepochsecondsbucket: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.fromEpochSecondsBucket}" + transform: null + udf: null +fromepochsecondsbucketmv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.fromEpochSecondsBucketMV}" + transform: null + udf: null +fromepochsecondsmv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.fromEpochSecondsMV}" + transform: null + udf: null +fromiso8601: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.fromIso8601}" + transform: null + udf: null +fromiso8601mv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.fromIso8601MV}" + transform: null + udf: null +fromtimestamp: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.fromTimestamp}" + transform: null + udf: null +fromtimestampmv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.fromTimestampMV}" + transform: null + udf: null +fromull: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.function.scalar.SketchFunctions.fromULL}" + transform: null + udf: null +fromutf8: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.fromUtf8}" + transform: null + udf: null +fromuuidbytes: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.fromUUIDBytes}" + transform: null + udf: null +gcd: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.gcd}" + transform: null + udf: null +generatedoublearray: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.generateDoubleArray}" + transform: null + udf: null +generatefloatarray: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.generateFloatArray}" + transform: null + udf: null +generateintarray: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.generateIntArray}" + transform: null + udf: null +generatelongarray: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.generateLongArray}" + transform: null + udf: null +geotoh3: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.core.geospatial.transform.function.ScalarFunctions.geoToH3,\ + \ 3: org.apache.pinot.core.geospatial.transform.function.ScalarFunctions.geoToH3]}" + transform: "org.apache.pinot.core.geospatial.transform.function.GeoToH3Function" + udf: null +getcpcsketchestimate: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.function.scalar.SketchFunctions.getCpcSketchEstimate}" + transform: null + udf: null +getinttuplesketchestimate: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.function.scalar.SketchFunctions.getIntTupleSketchEstimate}" + transform: null + udf: null +getthetasketchestimate: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.function.scalar.SketchFunctions.getThetaSketchEstimate}" + transform: null + udf: null +greaterthan: + scalar: "org.apache.pinot.common.function.scalar.comparison.GreaterThanScalarFunction" + transform: "org.apache.pinot.core.operator.transform.function.GreaterThanTransformFunction" + udf: null +greaterthanorequal: + scalar: "org.apache.pinot.common.function.scalar.comparison.GreaterThanOrEqualScalarFunction" + transform: "org.apache.pinot.core.operator.transform.function.GreaterThanOrEqualTransformFunction" + udf: null +greatest: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.greatest}" + transform: "org.apache.pinot.core.operator.transform.function.GreatestTransformFunction" + udf: null +griddisk: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.geospatial.transform.function.ScalarFunctions.gridDisk}" + transform: "org.apache.pinot.core.geospatial.transform.function.GridDiskFunction" + udf: null +griddistance: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.geospatial.transform.function.ScalarFunctions.gridDistance}" + transform: "org.apache.pinot.core.geospatial.transform.function.GridDistanceFunction" + udf: null +groovy: + scalar: null + transform: "org.apache.pinot.core.operator.transform.function.GroovyTransformFunction" + udf: null +hammingdistance: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.hammingDistance}" + transform: null + udf: null +hexdecimaltolong: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DataTypeConversionFunctions.hexDecimalToLong}" + transform: null + udf: null +hextobytes: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DataTypeConversionFunctions.hexToBytes}" + transform: null + udf: null +hour: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.hour,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.hour]}" + transform: "org.apache.pinot.core.operator.transform.function.DateTimeTransformFunction.Hour" + udf: null +hourmv: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.hourMV,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.hourMV]}" + transform: null + udf: null +hypot: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.hypot}" + transform: null + udf: null +ifnotfinite: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.ifNotFinite}" + transform: null + udf: null +in: + scalar: null + transform: "org.apache.pinot.core.operator.transform.function.InTransformFunction" + udf: null +inidset: + scalar: null + transform: "org.apache.pinot.core.operator.transform.function.InIdSetTransformFunction" + udf: null +innerproduct: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.VectorFunctions.innerProduct}" + transform: "org.apache.pinot.core.operator.transform.function.VectorTransformFunctions.InnerProductTransformFunction" + udf: null +intdiv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.intDiv}" + transform: null + udf: null +intdivorzero: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.intDivOrZero}" + transform: null + udf: null +intersectindices: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArrayFunctions.intersectIndices}" + transform: null + udf: null +intmaxtuplesketchintersect: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.function.scalar.SketchFunctions.intMaxTupleSketchIntersect}" + transform: null + udf: null +intmaxtuplesketchunion: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.core.function.scalar.SketchFunctions.intMaxTupleSketchUnion,\ + \ 3: org.apache.pinot.core.function.scalar.SketchFunctions.intMaxTupleSketchUnion]}" + transform: null + udf: null +intmintuplesketchintersect: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.function.scalar.SketchFunctions.intMinTupleSketchIntersect}" + transform: null + udf: null +intmintuplesketchunion: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.core.function.scalar.SketchFunctions.intMinTupleSketchUnion,\ + \ 3: org.apache.pinot.core.function.scalar.SketchFunctions.intMinTupleSketchUnion]}" + transform: null + udf: null +intsumtuplesketchdiff: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.function.scalar.SketchFunctions.intSumTupleSketchDiff}" + transform: null + udf: null +intsumtuplesketchintersect: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.function.scalar.SketchFunctions.intSumTupleSketchIntersect}" + transform: null + udf: null +intsumtuplesketchunion: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.core.function.scalar.SketchFunctions.intSumTupleSketchUnion,\ + \ 3: org.apache.pinot.core.function.scalar.SketchFunctions.intSumTupleSketchUnion]}" + transform: null + udf: null +isdistinctfrom: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ObjectFunctions.isDistinctFrom}" + transform: "org.apache.pinot.core.operator.transform.function.IsDistinctFromTransformFunction" + udf: null +isfalse: + scalar: null + transform: "org.apache.pinot.core.operator.transform.function.IsFalseTransformFunction" + udf: "org.apache.pinot.query.runtime.function.IsFalseUdf" +isfinite: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.isFinite}" + transform: null + udf: null +isinfinite: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.isInfinite}" + transform: null + udf: null +isjson: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.isJson}" + transform: null + udf: null +isnan: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.isNaN}" + transform: null + udf: null +isnotdistinctfrom: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ObjectFunctions.isNotDistinctFrom}" + transform: "org.apache.pinot.core.operator.transform.function.IsNotDistinctFromTransformFunction" + udf: null +isnotfalse: + scalar: null + transform: "org.apache.pinot.core.operator.transform.function.IsNotFalseTransformFunction" + udf: null +isnotnull: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ObjectFunctions.isNotNull}" + transform: "org.apache.pinot.core.operator.transform.function.IsNotNullTransformFunction" + udf: null +isnottrue: + scalar: null + transform: "org.apache.pinot.core.operator.transform.function.IsNotTrueTransformFunction" + udf: null +isnull: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ObjectFunctions.isNull}" + transform: "org.apache.pinot.core.operator.transform.function.IsNullTransformFunction" + udf: null +issubnetof: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.IpAddressFunctions.isSubnetOf}" + transform: null + udf: null +istrue: + scalar: null + transform: "org.apache.pinot.core.operator.transform.function.IsTrueTransformFunction" + udf: null +item: + scalar: null + transform: "org.apache.pinot.core.operator.transform.function.ItemTransformFunction" + udf: null +jsonextractindex: + scalar: null + transform: "org.apache.pinot.core.operator.transform.function.JsonExtractIndexTransformFunction" + udf: null +jsonextractkey: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.JsonFunctions.jsonExtractKey,\ + \ 3: org.apache.pinot.common.function.scalar.JsonFunctions.jsonExtractKey]}" + transform: "org.apache.pinot.core.operator.transform.function.JsonExtractKeyTransformFunction" + udf: null +jsonextractscalar: + scalar: null + transform: "org.apache.pinot.core.operator.transform.function.JsonExtractScalarTransformFunction" + udf: null +jsonformat: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.JsonFunctions.jsonFormat}" + transform: null + udf: null +jsonkeyvaluearraytomap: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.JsonFunctions.jsonKeyValueArrayToMap,\ + \ 3: org.apache.pinot.common.function.scalar.JsonFunctions.jsonKeyValueArrayToMap]}" + transform: null + udf: null +jsonpath: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.JsonFunctions.jsonPath}" + transform: null + udf: null +jsonpatharray: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.JsonFunctions.jsonPathArray}" + transform: null + udf: null +jsonpatharraydefaultempty: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.JsonFunctions.jsonPathArrayDefaultEmpty}" + transform: null + udf: null +jsonpathdouble: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.JsonFunctions.jsonPathDouble,\ + \ 3: org.apache.pinot.common.function.scalar.JsonFunctions.jsonPathDouble]}" + transform: null + udf: null +jsonpathexists: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.JsonFunctions.jsonPathExists}" + transform: null + udf: null +jsonpathlong: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.JsonFunctions.jsonPathLong,\ + \ 3: org.apache.pinot.common.function.scalar.JsonFunctions.jsonPathLong]}" + transform: null + udf: null +jsonpathstring: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.JsonFunctions.jsonPathString,\ + \ 3: org.apache.pinot.common.function.scalar.JsonFunctions.jsonPathString]}" + transform: null + udf: null +jsonstringtoarray: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.JsonFunctions.jsonStringToArray}" + transform: null + udf: null +jsonstringtolistormap: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.JsonFunctions.jsonStringToListOrMap}" + transform: null + udf: null +jsonstringtomap: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.JsonFunctions.jsonStringToMap}" + transform: null + udf: null +l1distance: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.VectorFunctions.l1Distance}" + transform: "org.apache.pinot.core.operator.transform.function.VectorTransformFunctions.L1DistanceTransformFunction" + udf: null +l2distance: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.VectorFunctions.l2Distance}" + transform: "org.apache.pinot.core.operator.transform.function.VectorTransformFunctions.L2DistanceTransformFunction" + udf: null +lcm: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.lcm}" + transform: null + udf: null +least: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.least}" + transform: "org.apache.pinot.core.operator.transform.function.LeastTransformFunction" + udf: null +left: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.leftSubStr}" + transform: null + udf: null +leftsubstr: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.leftSubStr}" + transform: null + udf: null +length: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.length}" + transform: null + udf: null +lessthan: + scalar: "org.apache.pinot.common.function.scalar.comparison.LessThanScalarFunction" + transform: "org.apache.pinot.core.operator.transform.function.LessThanTransformFunction" + udf: null +lessthanorequal: + scalar: "org.apache.pinot.common.function.scalar.comparison.LessThanOrEqualScalarFunction" + transform: "org.apache.pinot.core.operator.transform.function.LessThanOrEqualTransformFunction" + udf: null +like: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.regexp.RegexpLikeConstFunctions.like}" + transform: null + udf: null +likevar: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.regexp.RegexpLikeVarFunctions.likeVar}" + transform: null + udf: null +ln: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.ln}" + transform: "org.apache.pinot.core.operator.transform.function.SingleParamMathTransformFunction.LnTransformFunction" + udf: null +log: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.ln}" + transform: "org.apache.pinot.core.operator.transform.function.SingleParamMathTransformFunction.LnTransformFunction" + udf: null +log10: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.log10}" + transform: "org.apache.pinot.core.operator.transform.function.SingleParamMathTransformFunction.Log10TransformFunction" + udf: null +log2: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.log2}" + transform: "org.apache.pinot.core.operator.transform.function.SingleParamMathTransformFunction.Log2TransformFunction" + udf: null +longtohexdecimal: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DataTypeConversionFunctions.longToHexDecimal}" + transform: null + udf: null +lookup: + scalar: null + transform: "org.apache.pinot.core.operator.transform.function.LookupTransformFunction" + udf: null +lower: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.lower}" + transform: null + udf: null +lpad: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.lpad}" + transform: null + udf: null +ltrim: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.string.LTrimFunction.ltrim}" + transform: null + udf: null +mapvalue: + scalar: null + transform: "org.apache.pinot.core.operator.transform.function.MapValueTransformFunction" + udf: null +max: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.max}" + transform: null + udf: null +md2: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.HashFunctions.md2}" + transform: null + udf: null +md5: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.HashFunctions.md5}" + transform: null + udf: null +millisecond: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.millisecond,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.millisecond]}" + transform: "org.apache.pinot.core.operator.transform.function.DateTimeTransformFunction.Millisecond" + udf: null +millisecondmv: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.millisecondMV,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.millisecondMV]}" + transform: null + udf: null +min: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.min}" + transform: null + udf: null +minus: + scalar: "org.apache.pinot.common.function.scalar.arithmetic.MinusScalarFunction" + transform: "org.apache.pinot.core.operator.transform.function.SubtractionTransformFunction" + udf: null +minute: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.minute,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.minute]}" + transform: "org.apache.pinot.core.operator.transform.function.DateTimeTransformFunction.Minute" + udf: null +minutemv: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.minuteMV,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.minuteMV]}" + transform: null + udf: null +mod: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.mod}" + transform: "org.apache.pinot.core.operator.transform.function.ModuloTransformFunction" + udf: null +moduloorzero: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.moduloOrZero}" + transform: null + udf: null +month: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.month,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.month]}" + transform: "org.apache.pinot.core.operator.transform.function.DateTimeTransformFunction.Month" + udf: null +monthmv: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.monthMV,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.monthMV]}" + transform: null + udf: null +monthofyear: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.month,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.month]}" + transform: "org.apache.pinot.core.operator.transform.function.DateTimeTransformFunction.Month" + udf: null +monthofyearmv: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.monthMV,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.monthMV]}" + transform: null + udf: null +mult: + scalar: "org.apache.pinot.common.function.scalar.arithmetic.MultScalarFunction" + transform: "org.apache.pinot.core.operator.transform.function.MultiplicationTransformFunction" + udf: "org.apache.pinot.query.runtime.function.MultUdf" +murmurhash2: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.HashFunctions.murmurHash2}" + transform: null + udf: null +murmurhash2bit64: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.HashFunctions.murmurHash2Bit64,\ + \ 2: org.apache.pinot.common.function.scalar.HashFunctions.murmurHash2Bit64]}" + transform: null + udf: null +murmurhash2utf8: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.HashFunctions.murmurHash2UTF8}" + transform: null + udf: null +murmurhash3bit128: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.HashFunctions.murmurHash3Bit128}" + transform: null + udf: null +murmurhash3bit32: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.HashFunctions.murmurHash3Bit32}" + transform: null + udf: null +murmurhash3bit64: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.HashFunctions.murmurHash3Bit64}" + transform: null + udf: null +murmurhash3x64bit128: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.HashFunctions.murmurHash3X64Bit128}" + transform: null + udf: null +murmurhash3x64bit32: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.HashFunctions.murmurHash3X64Bit32}" + transform: null + udf: null +murmurhash3x64bit64: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.HashFunctions.murmurHash3X64Bit64}" + transform: null + udf: null +negate: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.negate}" + transform: null + udf: null +normalize: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.StringFunctions.normalize,\ + \ 2: org.apache.pinot.common.function.scalar.StringFunctions.normalize]}" + transform: null + udf: null +not: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.LogicalFunctions.not}" + transform: "org.apache.pinot.core.operator.transform.function.NotOperatorTransformFunction" + udf: "org.apache.pinot.query.runtime.function.NotUdf" +notequals: + scalar: "org.apache.pinot.common.function.scalar.comparison.NotEqualsScalarFunction" + transform: "org.apache.pinot.core.operator.transform.function.NotEqualsTransformFunction" + udf: null +notin: + scalar: null + transform: "org.apache.pinot.core.operator.transform.function.NotInTransformFunction" + udf: null +now: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.now}" + transform: null + udf: null +nullif: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ObjectFunctions.nullIf}" + transform: null + udf: null +or: + scalar: null + transform: "org.apache.pinot.core.operator.transform.function.OrOperatorTransformFunction" + udf: "org.apache.pinot.query.runtime.function.OrUdf" +plus: + scalar: "org.apache.pinot.common.function.scalar.arithmetic.PlusScalarFunction" + transform: "org.apache.pinot.core.operator.transform.function.AdditionTransformFunction" + udf: "org.apache.pinot.query.runtime.function.PlusUdf" +positivemodulo: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.positiveModulo}" + transform: null + udf: null +pow: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.power}" + transform: "org.apache.pinot.core.operator.transform.function.PowerTransformFunction" + udf: null +power: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.power}" + transform: "org.apache.pinot.core.operator.transform.function.PowerTransformFunction" + udf: null +prefixes: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.prefixes}" + transform: null + udf: null +prefixeswithprefix: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.prefixesWithPrefix}" + transform: null + udf: null +quarter: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.quarter,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.quarter]}" + transform: "org.apache.pinot.core.operator.transform.function.DateTimeTransformFunction.Quarter" + udf: null +quartermv: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.quarterMV,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.quarterMV]}" + transform: null + udf: null +queryengine: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.InternalFunctions.queryEngine}" + transform: null + udf: null +radians: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.TrigonometricFunctions.radians}" + transform: "org.apache.pinot.core.operator.transform.function.TrigonometricTransformFunctions.RadiansTransformFunction" + udf: null +regexpextract: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.regexp.RegexpExtractConstFunctions.regexpExtract,\ + \ 3: org.apache.pinot.common.function.scalar.regexp.RegexpExtractConstFunctions.regexpExtract,\ + \ 4: org.apache.pinot.common.function.scalar.regexp.RegexpExtractConstFunctions.regexpExtract]}" + transform: "org.apache.pinot.core.operator.transform.function.RegexpExtractTransformFunction" + udf: null +regexpextractvar: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.regexp.RegexpExtractVarFunctions.regexpExtractVar,\ + \ 3: org.apache.pinot.common.function.scalar.regexp.RegexpExtractVarFunctions.regexpExtractVar,\ + \ 4: org.apache.pinot.common.function.scalar.regexp.RegexpExtractVarFunctions.regexpExtractVar]}" + transform: null + udf: null +regexplike: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.regexp.RegexpLikeConstFunctions.regexpLike,\ + \ 3: org.apache.pinot.common.function.scalar.regexp.RegexpLikeConstFunctions.regexpLike]}" + transform: null + udf: null +regexplikevar: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.regexp.RegexpLikeVarFunctions.regexpLikeVar,\ + \ 3: org.apache.pinot.common.function.scalar.regexp.RegexpLikeVarFunctions.regexpLikeVar]}" + transform: null + udf: null +regexpreplace: + scalar: "ArgumentCountBasedScalarFunction{[3: org.apache.pinot.common.function.scalar.regexp.RegexpReplaceConstFunctions.regexpReplace,\ + \ 4: org.apache.pinot.common.function.scalar.regexp.RegexpReplaceConstFunctions.regexpReplace,\ + \ 5: org.apache.pinot.common.function.scalar.regexp.RegexpReplaceConstFunctions.regexpReplace,\ + \ 6: org.apache.pinot.common.function.scalar.regexp.RegexpReplaceConstFunctions.regexpReplace]}" + transform: null + udf: null +regexpreplacevar: + scalar: "ArgumentCountBasedScalarFunction{[3: org.apache.pinot.common.function.scalar.regexp.RegexpReplaceVarFunctions.regexpReplaceVar,\ + \ 4: org.apache.pinot.common.function.scalar.regexp.RegexpReplaceVarFunctions.regexpReplaceVar,\ + \ 5: org.apache.pinot.common.function.scalar.regexp.RegexpReplaceVarFunctions.regexpReplaceVar,\ + \ 6: org.apache.pinot.common.function.scalar.regexp.RegexpReplaceVarFunctions.regexpReplaceVar]}" + transform: null + udf: null +remove: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.remove}" + transform: null + udf: null +repeat: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.StringFunctions.repeat,\ + \ 3: org.apache.pinot.common.function.scalar.StringFunctions.repeat]}" + transform: null + udf: null +replace: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.string.StringFunctions.replace}" + transform: null + udf: null +reqid: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.InternalFunctions.reqId}" + transform: null + udf: null +reverse: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.reverse}" + transform: null + udf: null +right: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.rightSubStr}" + transform: null + udf: null +rightsubstr: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.rightSubStr}" + transform: null + udf: null +round: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.round}" + transform: null + udf: null +rounddecimal: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.ArithmeticFunctions.roundDecimal,\ + \ 2: org.apache.pinot.common.function.scalar.ArithmeticFunctions.roundDecimal]}" + transform: "org.apache.pinot.core.operator.transform.function.RoundDecimalTransformFunction" + udf: null +roundmv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.roundMV}" + transform: null + udf: null +rpad: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.rpad}" + transform: null + udf: null +rtrim: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.string.RTrimFunction.rtrim}" + transform: null + udf: null +second: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.second,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.second]}" + transform: "org.apache.pinot.core.operator.transform.function.DateTimeTransformFunction.Second" + udf: null +secondmv: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.secondMV,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.secondMV]}" + transform: null + udf: null +sha: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.HashFunctions.sha}" + transform: null + udf: null +sha224: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.HashFunctions.sha224}" + transform: null + udf: null +sha256: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.HashFunctions.sha256}" + transform: null + udf: null +sha512: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.HashFunctions.sha512}" + transform: null + udf: null +sign: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.sign}" + transform: "org.apache.pinot.core.operator.transform.function.SingleParamMathTransformFunction.SignTransformFunction" + udf: null +sin: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.TrigonometricFunctions.sin}" + transform: "org.apache.pinot.core.operator.transform.function.TrigonometricTransformFunctions.SinTransformFunction" + udf: null +sinh: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.TrigonometricFunctions.sinh}" + transform: "org.apache.pinot.core.operator.transform.function.TrigonometricTransformFunctions.SinhTransformFunction" + udf: null +sleep: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.sleep}" + transform: null + udf: null +split: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.StringFunctions.split,\ + \ 3: org.apache.pinot.common.function.scalar.StringFunctions.split]}" + transform: null + udf: null +splitpart: + scalar: "ArgumentCountBasedScalarFunction{[3: org.apache.pinot.common.function.scalar.StringFunctions.splitPart,\ + \ 4: org.apache.pinot.common.function.scalar.StringFunctions.splitPart]}" + transform: null + udf: null +sqrt: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.ArithmeticFunctions.sqrt}" + transform: "org.apache.pinot.core.operator.transform.function.SingleParamMathTransformFunction.SqrtTransformFunction" + udf: null +stageid: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.query.function.InternalMseFunctions.stageId}" + transform: null + udf: null +starea: + scalar: null + transform: "org.apache.pinot.core.geospatial.transform.function.StAreaFunction" + udf: null +startswith: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.startsWith}" + transform: null + udf: null +starttime: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.InternalFunctions.startTime}" + transform: null + udf: null +stasbinary: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.geospatial.transform.function.ScalarFunctions.stAsBinary}" + transform: "org.apache.pinot.core.geospatial.transform.function.StAsBinaryFunction" + udf: null +stasgeojson: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.geospatial.transform.function.ScalarFunctions.stAsGeoJson}" + transform: "org.apache.pinot.core.geospatial.transform.function.StAsGeoJsonFunction" + udf: null +stastext: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.geospatial.transform.function.ScalarFunctions.stAsText}" + transform: "org.apache.pinot.core.geospatial.transform.function.StAsTextFunction" + udf: null +stcontains: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.geospatial.transform.function.ScalarFunctions.stContains}" + transform: "org.apache.pinot.core.geospatial.transform.function.StContainsFunction" + udf: null +stdistance: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.geospatial.transform.function.ScalarFunctions.stDistance}" + transform: "org.apache.pinot.core.geospatial.transform.function.StDistanceFunction" + udf: null +stequals: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.geospatial.transform.function.ScalarFunctions.stEquals}" + transform: "org.apache.pinot.core.geospatial.transform.function.StEqualsFunction" + udf: null +stgeogfromgeojson: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.geospatial.transform.function.ScalarFunctions.stGeogFromGeoJson}" + transform: "org.apache.pinot.core.geospatial.transform.function.StGeogFromGeoJsonFunction" + udf: null +stgeogfromtext: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.geospatial.transform.function.ScalarFunctions.stGeogFromText}" + transform: "org.apache.pinot.core.geospatial.transform.function.StGeogFromTextFunction" + udf: null +stgeogfromwkb: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.geospatial.transform.function.ScalarFunctions.stGeogFromWKB}" + transform: "org.apache.pinot.core.geospatial.transform.function.StGeogFromWKBFunction" + udf: null +stgeometrytype: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.geospatial.transform.function.ScalarFunctions.stGeometryType}" + transform: "org.apache.pinot.core.geospatial.transform.function.StGeometryTypeFunction" + udf: null +stgeomfromgeojson: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.geospatial.transform.function.ScalarFunctions.stGeomFromGeoJson}" + transform: "org.apache.pinot.core.geospatial.transform.function.StGeomFromGeoJsonFunction" + udf: null +stgeomfromtext: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.geospatial.transform.function.ScalarFunctions.stGeomFromText}" + transform: "org.apache.pinot.core.geospatial.transform.function.StGeomFromTextFunction" + udf: null +stgeomfromwkb: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.geospatial.transform.function.ScalarFunctions.stGeomFromWKB}" + transform: "org.apache.pinot.core.geospatial.transform.function.StGeomFromWKBFunction" + udf: null +stpoint: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.core.geospatial.transform.function.ScalarFunctions.stPoint,\ + \ 3: org.apache.pinot.core.geospatial.transform.function.ScalarFunctions.stPoint]}" + transform: "org.apache.pinot.core.geospatial.transform.function.StPointFunction" + udf: null +stpolygon: + scalar: null + transform: "org.apache.pinot.core.geospatial.transform.function.StPolygonFunction" + udf: null +strcmp: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.strcmp}" + transform: null + udf: null +stringtoarray: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.StringFunctions.split,\ + \ 3: org.apache.pinot.common.function.scalar.StringFunctions.split]}" + transform: null + udf: null +strpos: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.StringFunctions.strpos,\ + \ 3: org.apache.pinot.common.function.scalar.StringFunctions.strpos]}" + transform: null + udf: null +strrpos: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.StringFunctions.strrpos,\ + \ 3: org.apache.pinot.common.function.scalar.StringFunctions.strrpos]}" + transform: null + udf: null +stwithin: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.geospatial.transform.function.ScalarFunctions.stWithin}" + transform: "org.apache.pinot.core.geospatial.transform.function.StWithinFunction" + udf: null +sub: + scalar: "org.apache.pinot.common.function.scalar.arithmetic.MinusScalarFunction" + transform: "org.apache.pinot.core.operator.transform.function.SubtractionTransformFunction" + udf: null +substr: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.StringFunctions.substr,\ + \ 3: org.apache.pinot.common.function.scalar.StringFunctions.substr]}" + transform: null + udf: null +substring: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.StringFunctions.substring,\ + \ 3: org.apache.pinot.common.function.scalar.StringFunctions.substring]}" + transform: null + udf: null +suffixes: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.suffixes}" + transform: null + udf: null +suffixeswithsuffix: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.suffixesWithSuffix}" + transform: null + udf: null +tan: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.TrigonometricFunctions.tan}" + transform: "org.apache.pinot.core.operator.transform.function.TrigonometricTransformFunctions.TanTransformFunction" + udf: null +tanh: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.TrigonometricFunctions.tanh}" + transform: "org.apache.pinot.core.operator.transform.function.TrigonometricTransformFunctions.TanhTransformFunction" + udf: null +testfunc1: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.function.FunctionRegistryTest.testScalarFunction}" + transform: null + udf: null +testfunc2: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.function.FunctionRegistryTest.testScalarFunction}" + transform: null + udf: null +textmatch: + scalar: null + transform: "org.apache.pinot.core.operator.transform.function.TextMatchTransformFunction" + udf: null +thetasketchdiff: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.function.scalar.SketchFunctions.thetaSketchDiff}" + transform: null + udf: null +thetasketchintersect: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.core.function.scalar.SketchFunctions.thetaSketchIntersect,\ + \ 3: org.apache.pinot.core.function.scalar.SketchFunctions.thetaSketchIntersect,\ + \ 4: org.apache.pinot.core.function.scalar.SketchFunctions.thetaSketchIntersect,\ + \ 5: org.apache.pinot.core.function.scalar.SketchFunctions.thetaSketchIntersect]}" + transform: null + udf: null +thetasketchtostring: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.function.scalar.SketchFunctions.thetaSketchToString}" + transform: null + udf: null +thetasketchunion: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.core.function.scalar.SketchFunctions.thetaSketchUnion,\ + \ 3: org.apache.pinot.core.function.scalar.SketchFunctions.thetaSketchUnion, 4:\ + \ org.apache.pinot.core.function.scalar.SketchFunctions.thetaSketchUnion, 5: org.apache.pinot.core.function.scalar.SketchFunctions.thetaSketchUnion]}" + transform: null + udf: null +timeconvert: + scalar: null + transform: "org.apache.pinot.core.operator.transform.function.TimeConversionTransformFunction" + udf: null +times: + scalar: "org.apache.pinot.common.function.scalar.arithmetic.MultScalarFunction" + transform: "org.apache.pinot.core.operator.transform.function.MultiplicationTransformFunction" + udf: "org.apache.pinot.query.runtime.function.MultUdf" +timeseriesbucket: + scalar: null + transform: "org.apache.pinot.core.operator.transform.function.TimeSeriesBucketTransformFunction" + udf: null +timestampadd: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.timestampAdd}" + transform: null + udf: null +timestampaddmv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.timestampAddMV}" + transform: null + udf: null +timestampdiff: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.timestampDiff}" + transform: null + udf: null +timestampdiffmv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.timestampDiffMV}" + transform: null + udf: null +timestampdiffmvreverse: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.timestampDiffMVReverse}" + transform: null + udf: null +timezonehour: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.timezoneHour,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.timezoneHour]}" + transform: null + udf: null +timezonehourmv: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.timezoneHourMV,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.timezoneHourMV]}" + transform: null + udf: null +timezoneminute: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.timezoneMinute,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.timezoneMinute]}" + transform: null + udf: null +timezoneminutemv: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.timezoneMinuteMV,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.timezoneMinuteMV]}" + transform: null + udf: null +toascii: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.toAscii}" + transform: null + udf: null +tobase64: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.toBase64}" + transform: null + udf: null +tobytes: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.toBytes}" + transform: null + udf: null +tocpcsketch: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.core.function.scalar.SketchFunctions.toCpcSketch,\ + \ 2: org.apache.pinot.core.function.scalar.SketchFunctions.toCpcSketch]}" + transform: null + udf: null +todatetime: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.DateTimeFunctions.toDateTime,\ + \ 3: org.apache.pinot.common.function.scalar.DateTimeFunctions.toDateTime]}" + transform: null + udf: "org.apache.pinot.query.runtime.function.ToDateTimeUdf" +todatetimemv: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.DateTimeFunctions.toDateTimeMV,\ + \ 3: org.apache.pinot.common.function.scalar.DateTimeFunctions.toDateTimeMV]}" + transform: null + udf: null +toepochdays: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toEpochDays}" + transform: null + udf: null +toepochdaysbucket: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toEpochDaysBucket}" + transform: null + udf: null +toepochdaysbucketmv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toEpochDaysBucketMV}" + transform: null + udf: null +toepochdaysmv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toEpochDaysMV}" + transform: null + udf: null +toepochdaysrounded: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toEpochDaysRounded}" + transform: null + udf: null +toepochdaysroundedmv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toEpochDaysRoundedMV}" + transform: null + udf: null +toepochhours: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toEpochHours}" + transform: null + udf: null +toepochhoursbucket: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toEpochHoursBucket}" + transform: null + udf: null +toepochhoursbucketmv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toEpochHoursBucketMV}" + transform: null + udf: null +toepochhoursmv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toEpochHoursMV}" + transform: null + udf: null +toepochhoursrounded: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toEpochHoursRounded}" + transform: null + udf: null +toepochhoursroundedmv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toEpochHoursRoundedMV}" + transform: null + udf: null +toepochminutes: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toEpochMinutes}" + transform: null + udf: null +toepochminutesbucket: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toEpochMinutesBucket}" + transform: null + udf: null +toepochminutesbucketmv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toEpochMinutesBucketMV}" + transform: null + udf: null +toepochminutesmv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toEpochMinutesMV}" + transform: null + udf: null +toepochminutesrounded: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toEpochMinutesRounded}" + transform: null + udf: null +toepochminutesroundedmv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toEpochMinutesRoundedMV}" + transform: null + udf: null +toepochseconds: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toEpochSeconds}" + transform: null + udf: null +toepochsecondsbucket: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toEpochSecondsBucket}" + transform: null + udf: null +toepochsecondsbucketmv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toEpochSecondsBucketMV}" + transform: null + udf: null +toepochsecondsmv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toEpochSecondsMV}" + transform: null + udf: null +toepochsecondsrounded: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toEpochSecondsRounded}" + transform: null + udf: null +toepochsecondsroundedmv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toEpochSecondsRoundedMV}" + transform: null + udf: null +togeometry: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.geospatial.transform.function.ScalarFunctions.toGeometry}" + transform: null + udf: null +tohll: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.core.function.scalar.SketchFunctions.toHLL,\ + \ 2: org.apache.pinot.core.function.scalar.SketchFunctions.toHLL]}" + transform: null + udf: null +tointegersumtuplesketch: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.core.function.scalar.SketchFunctions.toIntegerSumTupleSketch,\ + \ 3: org.apache.pinot.core.function.scalar.SketchFunctions.toIntegerSumTupleSketch]}" + transform: null + udf: null +toiso8601: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toIso8601}" + transform: null + udf: null +toiso8601mv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toIso8601MV}" + transform: null + udf: null +tojsonmapstr: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.JsonFunctions.toJsonMapStr}" + transform: null + udf: null +tosphericalgeography: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.core.geospatial.transform.function.ScalarFunctions.toSphericalGeography}" + transform: null + udf: null +tothetasketch: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.core.function.scalar.SketchFunctions.toThetaSketch,\ + \ 2: org.apache.pinot.core.function.scalar.SketchFunctions.toThetaSketch]}" + transform: null + udf: null +totimestamp: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toTimestamp}" + transform: null + udf: null +totimestampmv: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.DateTimeFunctions.toTimestampMV}" + transform: null + udf: null +toull: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.core.function.scalar.SketchFunctions.toULL,\ + \ 2: org.apache.pinot.core.function.scalar.SketchFunctions.toULL]}" + transform: null + udf: null +toutf8: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.toUtf8}" + transform: null + udf: null +touuidbytes: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.toUUIDBytes}" + transform: null + udf: null +trim: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.StringFunctions.trim,\ + \ 3: org.apache.pinot.common.function.scalar.StringFunctions.trim]}" + transform: null + udf: null +truncate: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.ArithmeticFunctions.truncate,\ + \ 2: org.apache.pinot.common.function.scalar.ArithmeticFunctions.truncate]}" + transform: "org.apache.pinot.core.operator.transform.function.TruncateDecimalTransformFunction" + udf: null +uniquengrams: + scalar: "ArgumentCountBasedScalarFunction{[2: org.apache.pinot.common.function.scalar.string.NgramFunctions.uniqueNgrams,\ + \ 3: org.apache.pinot.common.function.scalar.string.NgramFunctions.uniqueNgrams]}" + transform: null + udf: null +upper: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.StringFunctions.upper}" + transform: null + udf: null +urldecode: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.urlDecode}" + transform: null + udf: null +urldecodeformcomponent: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.urlDecodeFormComponent}" + transform: null + udf: null +urldomain: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.urlDomain}" + transform: null + udf: null +urldomainwithoutwww: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.urlDomainWithoutWWW}" + transform: null + udf: null +urlencode: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.urlEncode}" + transform: null + udf: null +urlencodeformcomponent: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.urlEncodeFormComponent}" + transform: null + udf: null +urlfirstsignificantsubdomain: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.urlFirstSignificantSubdomain}" + transform: null + udf: null +urlfragment: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.urlFragment}" + transform: null + udf: null +urlhierarchy: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.urlHierarchy}" + transform: null + udf: null +urlnetloc: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.urlNetloc}" + transform: null + udf: null +urlpath: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.urlPath}" + transform: null + udf: null +urlpathhierarchy: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.urlPathHierarchy}" + transform: null + udf: null +urlpathwithquery: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.urlPathWithQuery}" + transform: null + udf: null +urlport: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.urlPort}" + transform: null + udf: null +urlprotocol: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.urlProtocol}" + transform: null + udf: null +urlquerystring: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.urlQueryString}" + transform: null + udf: null +urlquerystringandfragment: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.urlQueryStringAndFragment}" + transform: null + udf: null +urltopleveldomain: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.UrlFunctions.urlTopLevelDomain}" + transform: null + udf: null +valuein: + scalar: null + transform: "org.apache.pinot.core.operator.transform.function.ValueInTransformFunction" + udf: null +vectordims: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.VectorFunctions.vectorDims}" + transform: "org.apache.pinot.core.operator.transform.function.VectorTransformFunctions.VectorDimsTransformFunction" + udf: null +vectornorm: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.common.function.scalar.VectorFunctions.vectorNorm}" + transform: "org.apache.pinot.core.operator.transform.function.VectorTransformFunctions.VectorNormTransformFunction" + udf: null +week: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.week,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.week]}" + transform: "org.apache.pinot.core.operator.transform.function.DateTimeTransformFunction.WeekOfYear" + udf: null +weekmv: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.weekMV,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.weekMV]}" + transform: null + udf: null +weekofyear: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.week,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.week]}" + transform: "org.apache.pinot.core.operator.transform.function.DateTimeTransformFunction.WeekOfYear" + udf: null +weekofyearmv: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.weekMV,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.weekMV]}" + transform: null + udf: null +workerid: + scalar: "ArgumentCountBasedScalarFunction{org.apache.pinot.query.function.InternalMseFunctions.workerId}" + transform: null + udf: null +year: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.year,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.year]}" + transform: "org.apache.pinot.core.operator.transform.function.DateTimeTransformFunction.Year" + udf: "org.apache.pinot.query.runtime.function.YearUdf" +yearmv: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.yearMV,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.yearMV]}" + transform: null + udf: null +yearofweek: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.yearOfWeek,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.yearOfWeek]}" + transform: "org.apache.pinot.core.operator.transform.function.DateTimeTransformFunction.YearOfWeek" + udf: null +yearofweekmv: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.yearOfWeekMV,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.yearOfWeekMV]}" + transform: null + udf: null +yow: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.yearOfWeek,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.yearOfWeek]}" + transform: "org.apache.pinot.core.operator.transform.function.DateTimeTransformFunction.YearOfWeek" + udf: null +yowmv: + scalar: "ArgumentCountBasedScalarFunction{[1: org.apache.pinot.common.function.scalar.DateTimeFunctions.yearOfWeekMV,\ + \ 2: org.apache.pinot.common.function.scalar.DateTimeFunctions.yearOfWeekMV]}" + transform: null + udf: null diff --git a/pinot-integration-tests/src/test/resources/udf-test-results/and.yaml b/pinot-integration-tests/src/test/resources/udf-test-results/and.yaml new file mode 100644 index 000000000000..b5907ae8d2c6 --- /dev/null +++ b/pinot-integration-tests/src/test/resources/udf-test-results/and.yaml @@ -0,0 +1,303 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# This file is auto-generated by the UDF test framework. Do not edit it manually. +# Use the org.apache.pinot.integration.tests.udfUdfTest.generateSnapshots() method to regenerate it. + +--- +Ingestion time transformer: + '(left: boolean, right: boolean) -> boolean': + entries: + false and false: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + false and true: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + null and null: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + null and true: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + true and false: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + true and null: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + true and true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +MSE intermediate stage (with null handling): + '(left: boolean, right: boolean) -> boolean': + entries: + false and false: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + false and true: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + null and null: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + null and true: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + true and false: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + true and null: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + true and true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +MSE intermediate stage (without null handling): + '(left: boolean, right: boolean) -> boolean': + entries: + false and false: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + false and true: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + null and null: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + null and true: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + true and false: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + true and null: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + true and true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE predicate (with null handling): + '(left: boolean, right: boolean) -> boolean': + entries: + false and false: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + false and true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null and null: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null and true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + true and false: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + true and null: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + true and true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE predicate (without null handling): + '(left: boolean, right: boolean) -> boolean': + entries: + false and false: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + false and true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null and null: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null and true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + true and false: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + true and null: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + true and true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE projection (with null handling): + '(left: boolean, right: boolean) -> boolean': + entries: + false and false: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + false and true: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + null and null: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + null and true: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + true and false: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + true and null: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + true and true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE projection (without null handling): + '(left: boolean, right: boolean) -> boolean': + entries: + false and false: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + false and true: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + null and null: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + null and true: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + true and false: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + true and null: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + true and true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null diff --git a/pinot-integration-tests/src/test/resources/udf-test-results/array.yaml b/pinot-integration-tests/src/test/resources/udf-test-results/array.yaml new file mode 100644 index 000000000000..9e92ba11d87a --- /dev/null +++ b/pinot-integration-tests/src/test/resources/udf-test-results/array.yaml @@ -0,0 +1,30 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# This file is auto-generated by the UDF test framework. Do not edit it manually. +# Use the org.apache.pinot.integration.tests.udfUdfTest.generateSnapshots() method to regenerate it. + +--- +Ingestion time transformer: {} +MSE intermediate stage (with null handling): {} +MSE intermediate stage (without null handling): {} +SSE predicate (with null handling): {} +SSE predicate (without null handling): {} +SSE projection (with null handling): {} +SSE projection (without null handling): {} diff --git a/pinot-integration-tests/src/test/resources/udf-test-results/arrayConcatDouble.yaml b/pinot-integration-tests/src/test/resources/udf-test-results/arrayConcatDouble.yaml new file mode 100644 index 000000000000..b88e9aa7c3c7 --- /dev/null +++ b/pinot-integration-tests/src/test/resources/udf-test-results/arrayConcatDouble.yaml @@ -0,0 +1,339 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# This file is auto-generated by the UDF test framework. Do not edit it manually. +# Use the org.apache.pinot.integration.tests.udfUdfTest.generateSnapshots() method to regenerate it. + +--- +Ingestion time transformer: + '(value1: ARRAY(double), value2: ARRAY(double)) -> ARRAY(double)': + entries: + empty with not empty: + actualResult: + - 1.0 + - 2.0 + equivalence: "EQUAL" + error: null + expectedResult: + - 1.0 + - 2.0 + not empty with empty: + actualResult: + - 1.0 + - 2.0 + equivalence: "EQUAL" + error: null + expectedResult: + - 1.0 + - 2.0 + not null concat null: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + null concat not null: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + null input: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + two empty arrays: + actualResult: [] + equivalence: "EQUAL" + error: null + expectedResult: [] + two not empty arrays: + actualResult: + - 1.0 + - 3.0 + - 2.0 + - 4.0 + equivalence: "EQUAL" + error: null + expectedResult: + - 1.0 + - 3.0 + - 2.0 + - 4.0 + error: false + errorMessage: null +MSE intermediate stage (with null handling): + '(value1: ARRAY(double), value2: ARRAY(double)) -> ARRAY(double)': + entries: + empty with not empty: + actualResult: [] + equivalence: null + error: "Unexpected value" + expectedResult: + - 1.0 + - 2.0 + not empty with empty: + actualResult: [] + equivalence: null + error: "Unexpected value" + expectedResult: + - 1.0 + - 2.0 + not null concat null: + actualResult: [] + equivalence: null + error: "Unexpected value" + expectedResult: null + null concat not null: + actualResult: [] + equivalence: null + error: "Unexpected value" + expectedResult: null + null input: + actualResult: [] + equivalence: null + error: "Unexpected value" + expectedResult: null + two empty arrays: + actualResult: [] + equivalence: "EQUAL" + error: null + expectedResult: [] + two not empty arrays: + actualResult: + - 1.0 + - 3.0 + - 2.0 + - 4.0 + equivalence: "EQUAL" + error: null + expectedResult: + - 1.0 + - 3.0 + - 2.0 + - 4.0 + error: false + errorMessage: null +MSE intermediate stage (without null handling): + '(value1: ARRAY(double), value2: ARRAY(double)) -> ARRAY(double)': + entries: + empty with not empty: + actualResult: + - -Infinity + - 1.0 + - 2.0 + equivalence: null + error: "Unexpected value" + expectedResult: + - 1.0 + - 2.0 + not empty with empty: + actualResult: + - 1.0 + - 2.0 + - -Infinity + equivalence: null + error: "Unexpected value" + expectedResult: + - 1.0 + - 2.0 + not null concat null: + actualResult: + - 1.0 + - 2.0 + - -Infinity + equivalence: null + error: "Unexpected value" + expectedResult: + - 1.0 + - 2.0 + null concat not null: + actualResult: + - -Infinity + - 1.0 + - 2.0 + equivalence: null + error: "Unexpected value" + expectedResult: + - 1.0 + - 2.0 + null input: + actualResult: + - -Infinity + - -Infinity + equivalence: null + error: "Unexpected value" + expectedResult: [] + two empty arrays: + actualResult: + - -Infinity + - -Infinity + equivalence: null + error: "Unexpected value" + expectedResult: [] + two not empty arrays: + actualResult: + - 1.0 + - 3.0 + - 2.0 + - 4.0 + equivalence: "EQUAL" + error: null + expectedResult: + - 1.0 + - 3.0 + - 2.0 + - 4.0 + error: false + errorMessage: null +SSE predicate (with null handling): + '(value1: ARRAY(double), value2: ARRAY(double)) -> ARRAY(double)': + entries: null + error: true + errorMessage: "\"Caught exception while initializing transform function: minus:\ + \ every argument of SUB transform function must be single-valued\"" +SSE predicate (without null handling): + '(value1: ARRAY(double), value2: ARRAY(double)) -> ARRAY(double)': + entries: null + error: true + errorMessage: "\"Caught exception while initializing transform function: minus:\ + \ every argument of SUB transform function must be single-valued\"" +SSE projection (with null handling): + '(value1: ARRAY(double), value2: ARRAY(double)) -> ARRAY(double)': + entries: + empty with not empty: + actualResult: [] + equivalence: null + error: "Unexpected value" + expectedResult: + - 1.0 + - 2.0 + not empty with empty: + actualResult: [] + equivalence: null + error: "Unexpected value" + expectedResult: + - 1.0 + - 2.0 + not null concat null: + actualResult: [] + equivalence: null + error: "Unexpected value" + expectedResult: null + null concat not null: + actualResult: [] + equivalence: null + error: "Unexpected value" + expectedResult: null + null input: + actualResult: [] + equivalence: null + error: "Unexpected value" + expectedResult: null + two empty arrays: + actualResult: [] + equivalence: "EQUAL" + error: null + expectedResult: [] + two not empty arrays: + actualResult: + - 1.0 + - 3.0 + - 2.0 + - 4.0 + equivalence: "EQUAL" + error: null + expectedResult: + - 1.0 + - 3.0 + - 2.0 + - 4.0 + error: false + errorMessage: null +SSE projection (without null handling): + '(value1: ARRAY(double), value2: ARRAY(double)) -> ARRAY(double)': + entries: + empty with not empty: + actualResult: + - -Infinity + - 1.0 + - 2.0 + equivalence: null + error: "Unexpected value" + expectedResult: + - 1.0 + - 2.0 + not empty with empty: + actualResult: + - 1.0 + - 2.0 + - -Infinity + equivalence: null + error: "Unexpected value" + expectedResult: + - 1.0 + - 2.0 + not null concat null: + actualResult: + - 1.0 + - 2.0 + - -Infinity + equivalence: null + error: "Unexpected value" + expectedResult: + - 1.0 + - 2.0 + null concat not null: + actualResult: + - -Infinity + - 1.0 + - 2.0 + equivalence: null + error: "Unexpected value" + expectedResult: + - 1.0 + - 2.0 + null input: + actualResult: + - -Infinity + - -Infinity + equivalence: null + error: "Unexpected value" + expectedResult: [] + two empty arrays: + actualResult: + - -Infinity + - -Infinity + equivalence: null + error: "Unexpected value" + expectedResult: [] + two not empty arrays: + actualResult: + - 1.0 + - 3.0 + - 2.0 + - 4.0 + equivalence: "EQUAL" + error: null + expectedResult: + - 1.0 + - 3.0 + - 2.0 + - 4.0 + error: false + errorMessage: null diff --git a/pinot-integration-tests/src/test/resources/udf-test-results/arrayElementAtInt.yaml b/pinot-integration-tests/src/test/resources/udf-test-results/arrayElementAtInt.yaml new file mode 100644 index 000000000000..4f8933212b9c --- /dev/null +++ b/pinot-integration-tests/src/test/resources/udf-test-results/arrayElementAtInt.yaml @@ -0,0 +1,233 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# This file is auto-generated by the UDF test framework. Do not edit it manually. +# Use the org.apache.pinot.integration.tests.udfUdfTest.generateSnapshots() method to regenerate it. + +--- +Ingestion time transformer: + '(arr: ARRAY(int), idx: int) -> int': + entries: + first element: + actualResult: 10 + equivalence: "EQUAL" + error: null + expectedResult: 10 + negative element: + actualResult: 0 + equivalence: "EQUAL" + error: null + expectedResult: 0 + out of bounds element: + actualResult: 0 + equivalence: "EQUAL" + error: null + expectedResult: 0 + second element: + actualResult: 20 + equivalence: "EQUAL" + error: null + expectedResult: 20 + zero element: + actualResult: 0 + equivalence: "EQUAL" + error: null + expectedResult: 0 + error: false + errorMessage: null +MSE intermediate stage (with null handling): + '(arr: ARRAY(int), idx: int) -> int': + entries: + first element: + actualResult: 10 + equivalence: "EQUAL" + error: null + expectedResult: 10 + negative element: + actualResult: 0 + equivalence: "EQUAL" + error: null + expectedResult: 0 + out of bounds element: + actualResult: 0 + equivalence: "EQUAL" + error: null + expectedResult: 0 + second element: + actualResult: 20 + equivalence: "EQUAL" + error: null + expectedResult: 20 + zero element: + actualResult: 0 + equivalence: "EQUAL" + error: null + expectedResult: 0 + error: false + errorMessage: null +MSE intermediate stage (without null handling): + '(arr: ARRAY(int), idx: int) -> int': + entries: + first element: + actualResult: 10 + equivalence: "EQUAL" + error: null + expectedResult: 10 + negative element: + actualResult: 0 + equivalence: "EQUAL" + error: null + expectedResult: 0 + out of bounds element: + actualResult: 0 + equivalence: "EQUAL" + error: null + expectedResult: 0 + second element: + actualResult: 20 + equivalence: "EQUAL" + error: null + expectedResult: 20 + zero element: + actualResult: 0 + equivalence: "EQUAL" + error: null + expectedResult: 0 + error: false + errorMessage: null +SSE predicate (with null handling): + '(arr: ARRAY(int), idx: int) -> int': + entries: + first element: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + negative element: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + out of bounds element: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + second element: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + zero element: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE predicate (without null handling): + '(arr: ARRAY(int), idx: int) -> int': + entries: + first element: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + negative element: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + out of bounds element: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + second element: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + zero element: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE projection (with null handling): + '(arr: ARRAY(int), idx: int) -> int': + entries: + first element: + actualResult: 10 + equivalence: "EQUAL" + error: null + expectedResult: 10 + negative element: + actualResult: 0 + equivalence: "EQUAL" + error: null + expectedResult: 0 + out of bounds element: + actualResult: 0 + equivalence: "EQUAL" + error: null + expectedResult: 0 + second element: + actualResult: 20 + equivalence: "EQUAL" + error: null + expectedResult: 20 + zero element: + actualResult: 0 + equivalence: "EQUAL" + error: null + expectedResult: 0 + error: false + errorMessage: null +SSE projection (without null handling): + '(arr: ARRAY(int), idx: int) -> int': + entries: + first element: + actualResult: 10 + equivalence: "EQUAL" + error: null + expectedResult: 10 + negative element: + actualResult: 0 + equivalence: "EQUAL" + error: null + expectedResult: 0 + out of bounds element: + actualResult: 0 + equivalence: "EQUAL" + error: null + expectedResult: 0 + second element: + actualResult: 20 + equivalence: "EQUAL" + error: null + expectedResult: 20 + zero element: + actualResult: 0 + equivalence: "EQUAL" + error: null + expectedResult: 0 + error: false + errorMessage: null diff --git a/pinot-integration-tests/src/test/resources/udf-test-results/arrayMax.yaml b/pinot-integration-tests/src/test/resources/udf-test-results/arrayMax.yaml new file mode 100644 index 000000000000..6a6c6720644e --- /dev/null +++ b/pinot-integration-tests/src/test/resources/udf-test-results/arrayMax.yaml @@ -0,0 +1,158 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# This file is auto-generated by the UDF test framework. Do not edit it manually. +# Use the org.apache.pinot.integration.tests.udfUdfTest.generateSnapshots() method to regenerate it. + +--- +Ingestion time transformer: + '(arr: ARRAY(int)) -> int': + entries: null + error: true + errorMessage: "Unsupported" +MSE intermediate stage (with null handling): + '(arr: ARRAY(int)) -> int': + entries: null + error: true + errorMessage: "Unsupported" +MSE intermediate stage (without null handling): + '(arr: ARRAY(int)) -> int': + entries: null + error: true + errorMessage: "Unsupported" +SSE predicate (with null handling): + '(arr: ARRAY(int)) -> int': + entries: + empty array: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + negative values: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + normal array: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null array: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + single element: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE predicate (without null handling): + '(arr: ARRAY(int)) -> int': + entries: + empty array: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + negative values: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + normal array: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null array: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + single element: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE projection (with null handling): + '(arr: ARRAY(int)) -> int': + entries: + empty array: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + negative values: + actualResult: 0 + equivalence: "EQUAL" + error: null + expectedResult: 0 + normal array: + actualResult: 3 + equivalence: "EQUAL" + error: null + expectedResult: 3 + null array: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + single element: + actualResult: 42 + equivalence: "EQUAL" + error: null + expectedResult: 42 + error: false + errorMessage: null +SSE projection (without null handling): + '(arr: ARRAY(int)) -> int': + entries: + empty array: + actualResult: -2147483648 + equivalence: "EQUAL" + error: null + expectedResult: -2147483648 + negative values: + actualResult: 0 + equivalence: "EQUAL" + error: null + expectedResult: 0 + normal array: + actualResult: 3 + equivalence: "EQUAL" + error: null + expectedResult: 3 + null array: + actualResult: -2147483648 + equivalence: "EQUAL" + error: null + expectedResult: -2147483648 + single element: + actualResult: 42 + equivalence: "EQUAL" + error: null + expectedResult: 42 + error: false + errorMessage: null diff --git a/pinot-integration-tests/src/test/resources/udf-test-results/arrayaverage.yaml b/pinot-integration-tests/src/test/resources/udf-test-results/arrayaverage.yaml new file mode 100644 index 000000000000..42b8a1a32052 --- /dev/null +++ b/pinot-integration-tests/src/test/resources/udf-test-results/arrayaverage.yaml @@ -0,0 +1,138 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# This file is auto-generated by the UDF test framework. Do not edit it manually. +# Use the org.apache.pinot.integration.tests.udfUdfTest.generateSnapshots() method to regenerate it. + +--- +Ingestion time transformer: + '(input: ARRAY(double)) -> double': + entries: null + error: true + errorMessage: "Unsupported" +MSE intermediate stage (with null handling): + '(input: ARRAY(double)) -> double': + entries: null + error: true + errorMessage: "Unsupported" +MSE intermediate stage (without null handling): + '(input: ARRAY(double)) -> double': + entries: null + error: true + errorMessage: "Unsupported" +SSE predicate (with null handling): + '(input: ARRAY(double)) -> double': + entries: + average of [1, 2, 3, 4]: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + average of [5]: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + empty array: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null input: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE predicate (without null handling): + '(input: ARRAY(double)) -> double': + entries: + average of [1, 2, 3, 4]: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + average of [5]: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + empty array: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null input: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE projection (with null handling): + '(input: ARRAY(double)) -> double': + entries: + average of [1, 2, 3, 4]: + actualResult: 2.5 + equivalence: "EQUAL" + error: null + expectedResult: 2.5 + average of [5]: + actualResult: 5.0 + equivalence: "EQUAL" + error: null + expectedResult: 5.0 + empty array: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + null input: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null +SSE projection (without null handling): + '(input: ARRAY(double)) -> double': + entries: + average of [1, 2, 3, 4]: + actualResult: 2.5 + equivalence: "EQUAL" + error: null + expectedResult: 2.5 + average of [5]: + actualResult: 5.0 + equivalence: "EQUAL" + error: null + expectedResult: 5.0 + empty array: + actualResult: -Infinity + equivalence: null + error: "Unexpected value" + expectedResult: null + null input: + actualResult: -Infinity + equivalence: "EQUAL" + error: null + expectedResult: -Infinity + error: false + errorMessage: null diff --git a/pinot-integration-tests/src/test/resources/udf-test-results/degrees.yaml b/pinot-integration-tests/src/test/resources/udf-test-results/degrees.yaml new file mode 100644 index 000000000000..be15e75c4367 --- /dev/null +++ b/pinot-integration-tests/src/test/resources/udf-test-results/degrees.yaml @@ -0,0 +1,233 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# This file is auto-generated by the UDF test framework. Do not edit it manually. +# Use the org.apache.pinot.integration.tests.udfUdfTest.generateSnapshots() method to regenerate it. + +--- +Ingestion time transformer: + '(radians: double) -> double': + entries: + degrees(0): + actualResult: 0.0 + equivalence: "EQUAL" + error: null + expectedResult: 0.0 + degrees(2*PI): + actualResult: 360.0 + equivalence: "EQUAL" + error: null + expectedResult: 360.0 + degrees(PI): + actualResult: 180.0 + equivalence: "EQUAL" + error: null + expectedResult: 180.0 + degrees(PI/2): + actualResult: 90.0 + equivalence: "EQUAL" + error: null + expectedResult: 90.0 + null input: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null +MSE intermediate stage (with null handling): + '(radians: double) -> double': + entries: + degrees(0): + actualResult: 0.0 + equivalence: "EQUAL" + error: null + expectedResult: 0.0 + degrees(2*PI): + actualResult: 360.0 + equivalence: "EQUAL" + error: null + expectedResult: 360.0 + degrees(PI): + actualResult: 180.0 + equivalence: "EQUAL" + error: null + expectedResult: 180.0 + degrees(PI/2): + actualResult: 90.0 + equivalence: "EQUAL" + error: null + expectedResult: 90.0 + null input: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null +MSE intermediate stage (without null handling): + '(radians: double) -> double': + entries: + degrees(0): + actualResult: 0.0 + equivalence: "EQUAL" + error: null + expectedResult: 0.0 + degrees(2*PI): + actualResult: 360.0 + equivalence: "EQUAL" + error: null + expectedResult: 360.0 + degrees(PI): + actualResult: 180.0 + equivalence: "EQUAL" + error: null + expectedResult: 180.0 + degrees(PI/2): + actualResult: 90.0 + equivalence: "EQUAL" + error: null + expectedResult: 90.0 + null input: + actualResult: 0.0 + equivalence: "EQUAL" + error: null + expectedResult: 0.0 + error: false + errorMessage: null +SSE predicate (with null handling): + '(radians: double) -> double': + entries: + degrees(0): + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + degrees(2*PI): + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + degrees(PI): + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + degrees(PI/2): + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null input: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE predicate (without null handling): + '(radians: double) -> double': + entries: + degrees(0): + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + degrees(2*PI): + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + degrees(PI): + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + degrees(PI/2): + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null input: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE projection (with null handling): + '(radians: double) -> double': + entries: + degrees(0): + actualResult: 0.0 + equivalence: "EQUAL" + error: null + expectedResult: 0.0 + degrees(2*PI): + actualResult: 360.0 + equivalence: "EQUAL" + error: null + expectedResult: 360.0 + degrees(PI): + actualResult: 180.0 + equivalence: "EQUAL" + error: null + expectedResult: 180.0 + degrees(PI/2): + actualResult: 90.0 + equivalence: "EQUAL" + error: null + expectedResult: 90.0 + null input: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null +SSE projection (without null handling): + '(radians: double) -> double': + entries: + degrees(0): + actualResult: 0.0 + equivalence: "EQUAL" + error: null + expectedResult: 0.0 + degrees(2*PI): + actualResult: 360.0 + equivalence: "EQUAL" + error: null + expectedResult: 360.0 + degrees(PI): + actualResult: 180.0 + equivalence: "EQUAL" + error: null + expectedResult: 180.0 + degrees(PI/2): + actualResult: 90.0 + equivalence: "EQUAL" + error: null + expectedResult: 90.0 + null input: + actualResult: 0.0 + equivalence: "EQUAL" + error: null + expectedResult: 0.0 + error: false + errorMessage: null diff --git a/pinot-integration-tests/src/test/resources/udf-test-results/isFalse.yaml b/pinot-integration-tests/src/test/resources/udf-test-results/isFalse.yaml new file mode 100644 index 000000000000..dde225bb7169 --- /dev/null +++ b/pinot-integration-tests/src/test/resources/udf-test-results/isFalse.yaml @@ -0,0 +1,144 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# This file is auto-generated by the UDF test framework. Do not edit it manually. +# Use the org.apache.pinot.integration.tests.udfUdfTest.generateSnapshots() method to regenerate it. + +--- +Ingestion time transformer: + '(input: int) -> boolean': + entries: null + error: true + errorMessage: "Unsupported" +MSE intermediate stage (with null handling): + '(input: int) -> boolean': + entries: null + error: true + errorMessage: "\"QueryValidationError: From line 15, column 3 to line 15, column\ + \ 19: No match found for function signature isfalse(). From line 15,\ + \ column 3 to line 15, column 19: No match found for function signature isfalse().\ + \ No match found for function signature isfalse()\"" +MSE intermediate stage (without null handling): + '(input: int) -> boolean': + entries: null + error: true + errorMessage: "\"QueryValidationError: From line 14, column 3 to line 14, column\ + \ 19: No match found for function signature isfalse(). From line 14,\ + \ column 3 to line 14, column 19: No match found for function signature isfalse().\ + \ No match found for function signature isfalse()\"" +SSE predicate (with null handling): + '(input: int) -> boolean': + entries: + input is -1 (not false): + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + input is 0 (false): + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + input is 1 (true): + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null input: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE predicate (without null handling): + '(input: int) -> boolean': + entries: + input is -1 (not false): + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + input is 0 (false): + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + input is 1 (true): + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null input: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE projection (with null handling): + '(input: int) -> boolean': + entries: + input is -1 (not false): + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + input is 0 (false): + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + input is 1 (true): + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + null input: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + error: false + errorMessage: null +SSE projection (without null handling): + '(input: int) -> boolean': + entries: + input is -1 (not false): + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + input is 0 (false): + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + input is 1 (true): + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + null input: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null diff --git a/pinot-integration-tests/src/test/resources/udf-test-results/mult.yaml b/pinot-integration-tests/src/test/resources/udf-test-results/mult.yaml new file mode 100644 index 000000000000..30c12db93a03 --- /dev/null +++ b/pinot-integration-tests/src/test/resources/udf-test-results/mult.yaml @@ -0,0 +1,520 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# This file is auto-generated by the UDF test framework. Do not edit it manually. +# Use the org.apache.pinot.integration.tests.udfUdfTest.generateSnapshots() method to regenerate it. + +--- +Ingestion time transformer: + '(arg0: big_decimal, arg1: big_decimal) -> big_decimal': + entries: + "2 * 3_big_decimal": + actualResult: 6.0 + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 6.0 + "2 * null_big_decimal": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null + '(arg0: double, arg1: double) -> double': + entries: + "2 * 3_double": + actualResult: 6.0 + equivalence: "EQUAL" + error: null + expectedResult: 6.0 + "2 * null_double": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null + '(arg0: float, arg1: float) -> float': + entries: + "2 * 3_float": + actualResult: 6.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 6.0 + "2 * null_float": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null + '(arg0: int, arg1: int) -> int': + entries: + "2 * 3_int": + actualResult: 6.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 6 + "2 * null_int": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null + '(arg0: long, arg1: long) -> long': + entries: + "2 * 3_long": + actualResult: 6.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 6 + "2 * null_long": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null +MSE intermediate stage (with null handling): + '(arg0: big_decimal, arg1: big_decimal) -> big_decimal': + entries: + "2 * 3_big_decimal": + actualResult: 6.0 + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 6.0 + "2 * null_big_decimal": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null + '(arg0: double, arg1: double) -> double': + entries: + "2 * 3_double": + actualResult: 6.0 + equivalence: "EQUAL" + error: null + expectedResult: 6.0 + "2 * null_double": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null + '(arg0: float, arg1: float) -> float': + entries: + "2 * 3_float": + actualResult: 6.0 + equivalence: "EQUAL" + error: null + expectedResult: 6.0 + "2 * null_float": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null + '(arg0: int, arg1: int) -> int': + entries: + "2 * 3_int": + actualResult: 6 + equivalence: "EQUAL" + error: null + expectedResult: 6 + "2 * null_int": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null + '(arg0: long, arg1: long) -> long': + entries: + "2 * 3_long": + actualResult: 6 + equivalence: "EQUAL" + error: null + expectedResult: 6 + "2 * null_long": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null +MSE intermediate stage (without null handling): + '(arg0: big_decimal, arg1: big_decimal) -> big_decimal': + entries: + "2 * 3_big_decimal": + actualResult: 6.0 + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 6.0 + "2 * null_big_decimal": + actualResult: 0.0 + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 0.0 + error: false + errorMessage: null + '(arg0: double, arg1: double) -> double': + entries: + "2 * 3_double": + actualResult: 6.0 + equivalence: "EQUAL" + error: null + expectedResult: 6.0 + "2 * null_double": + actualResult: 0.0 + equivalence: "EQUAL" + error: null + expectedResult: 0.0 + error: false + errorMessage: null + '(arg0: float, arg1: float) -> float': + entries: + "2 * 3_float": + actualResult: 6.0 + equivalence: "EQUAL" + error: null + expectedResult: 6.0 + "2 * null_float": + actualResult: 0.0 + equivalence: "EQUAL" + error: null + expectedResult: 0.0 + error: false + errorMessage: null + '(arg0: int, arg1: int) -> int': + entries: + "2 * 3_int": + actualResult: 6 + equivalence: "EQUAL" + error: null + expectedResult: 6 + "2 * null_int": + actualResult: 0 + equivalence: "EQUAL" + error: null + expectedResult: 0 + error: false + errorMessage: null + '(arg0: long, arg1: long) -> long': + entries: + "2 * 3_long": + actualResult: 6 + equivalence: "EQUAL" + error: null + expectedResult: 6 + "2 * null_long": + actualResult: 0 + equivalence: "EQUAL" + error: null + expectedResult: 0 + error: false + errorMessage: null +SSE predicate (with null handling): + '(arg0: big_decimal, arg1: big_decimal) -> big_decimal': + entries: + "2 * 3_big_decimal": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + "2 * null_big_decimal": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null + '(arg0: double, arg1: double) -> double': + entries: + "2 * 3_double": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + "2 * null_double": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null + '(arg0: float, arg1: float) -> float': + entries: + "2 * 3_float": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + "2 * null_float": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null + '(arg0: int, arg1: int) -> int': + entries: + "2 * 3_int": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + "2 * null_int": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null + '(arg0: long, arg1: long) -> long': + entries: + "2 * 3_long": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + "2 * null_long": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE predicate (without null handling): + '(arg0: big_decimal, arg1: big_decimal) -> big_decimal': + entries: + "2 * 3_big_decimal": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + "2 * null_big_decimal": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null + '(arg0: double, arg1: double) -> double': + entries: + "2 * 3_double": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + "2 * null_double": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null + '(arg0: float, arg1: float) -> float': + entries: + "2 * 3_float": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + "2 * null_float": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null + '(arg0: int, arg1: int) -> int': + entries: + "2 * 3_int": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + "2 * null_int": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null + '(arg0: long, arg1: long) -> long': + entries: + "2 * 3_long": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + "2 * null_long": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE projection (with null handling): + '(arg0: big_decimal, arg1: big_decimal) -> big_decimal': + entries: + "2 * 3_big_decimal": + actualResult: "6.0" + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 6.0 + "2 * null_big_decimal": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null + '(arg0: double, arg1: double) -> double': + entries: + "2 * 3_double": + actualResult: 6.0 + equivalence: "EQUAL" + error: null + expectedResult: 6.0 + "2 * null_double": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null + '(arg0: float, arg1: float) -> float': + entries: + "2 * 3_float": + actualResult: 6.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 6.0 + "2 * null_float": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null + '(arg0: int, arg1: int) -> int': + entries: + "2 * 3_int": + actualResult: 6.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 6 + "2 * null_int": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null + '(arg0: long, arg1: long) -> long': + entries: + "2 * 3_long": + actualResult: 6.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 6 + "2 * null_long": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null +SSE projection (without null handling): + '(arg0: big_decimal, arg1: big_decimal) -> big_decimal': + entries: + "2 * 3_big_decimal": + actualResult: "6.0" + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 6.0 + "2 * null_big_decimal": + actualResult: "0.0" + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 0.0 + error: false + errorMessage: null + '(arg0: double, arg1: double) -> double': + entries: + "2 * 3_double": + actualResult: 6.0 + equivalence: "EQUAL" + error: null + expectedResult: 6.0 + "2 * null_double": + actualResult: 0.0 + equivalence: "EQUAL" + error: null + expectedResult: 0.0 + error: false + errorMessage: null + '(arg0: float, arg1: float) -> float': + entries: + "2 * 3_float": + actualResult: 6.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 6.0 + "2 * null_float": + actualResult: 0.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 0.0 + error: false + errorMessage: null + '(arg0: int, arg1: int) -> int': + entries: + "2 * 3_int": + actualResult: 6.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 6 + "2 * null_int": + actualResult: 0.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 0 + error: false + errorMessage: null + '(arg0: long, arg1: long) -> long': + entries: + "2 * 3_long": + actualResult: 6.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 6 + "2 * null_long": + actualResult: 0.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 0 + error: false + errorMessage: null diff --git a/pinot-integration-tests/src/test/resources/udf-test-results/not.yaml b/pinot-integration-tests/src/test/resources/udf-test-results/not.yaml new file mode 100644 index 000000000000..290ee0862c9a --- /dev/null +++ b/pinot-integration-tests/src/test/resources/udf-test-results/not.yaml @@ -0,0 +1,163 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# This file is auto-generated by the UDF test framework. Do not edit it manually. +# Use the org.apache.pinot.integration.tests.udfUdfTest.generateSnapshots() method to regenerate it. + +--- +Ingestion time transformer: + '(value: boolean) -> boolean': + entries: + not false: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + not null: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + not true: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + error: false + errorMessage: null +MSE intermediate stage (with null handling): + '(value: boolean) -> boolean': + entries: + not false: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + not null: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + not true: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + error: false + errorMessage: null +MSE intermediate stage (without null handling): + '(value: boolean) -> boolean': + entries: + not false: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + not null: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + not true: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + error: false + errorMessage: null +SSE predicate (with null handling): + '(value: boolean) -> boolean': + entries: + not false: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + not null: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + not true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE predicate (without null handling): + '(value: boolean) -> boolean': + entries: + not false: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + not null: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + not true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE projection (with null handling): + '(value: boolean) -> boolean': + entries: + not false: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + not null: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + not true: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + error: false + errorMessage: null +SSE projection (without null handling): + '(value: boolean) -> boolean': + entries: + not false: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + not null: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + not true: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + error: false + errorMessage: null diff --git a/pinot-integration-tests/src/test/resources/udf-test-results/or.yaml b/pinot-integration-tests/src/test/resources/udf-test-results/or.yaml new file mode 100644 index 000000000000..7b44ddfe45e4 --- /dev/null +++ b/pinot-integration-tests/src/test/resources/udf-test-results/or.yaml @@ -0,0 +1,373 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# This file is auto-generated by the UDF test framework. Do not edit it manually. +# Use the org.apache.pinot.integration.tests.udfUdfTest.generateSnapshots() method to regenerate it. + +--- +Ingestion time transformer: + '(left: boolean, right: boolean) -> boolean': + entries: + false or false: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + false or null: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + false or true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null or false: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + null or null: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + null or true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + true or false: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + true or null: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + true or true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +MSE intermediate stage (with null handling): + '(left: boolean, right: boolean) -> boolean': + entries: + false or false: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + false or null: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + false or true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null or false: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + null or null: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + null or true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + true or false: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + true or null: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + true or true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +MSE intermediate stage (without null handling): + '(left: boolean, right: boolean) -> boolean': + entries: + false or false: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + false or null: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + false or true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null or false: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + null or null: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + null or true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + true or false: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + true or null: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + true or true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE predicate (with null handling): + '(left: boolean, right: boolean) -> boolean': + entries: + false or false: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + false or null: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + false or true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null or false: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null or null: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null or true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + true or false: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + true or null: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + true or true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE predicate (without null handling): + '(left: boolean, right: boolean) -> boolean': + entries: + false or false: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + false or null: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + false or true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null or false: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null or null: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null or true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + true or false: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + true or null: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + true or true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE projection (with null handling): + '(left: boolean, right: boolean) -> boolean': + entries: + false or false: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + false or null: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + false or true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null or false: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + null or null: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + null or true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + true or false: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + true or null: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + true or true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE projection (without null handling): + '(left: boolean, right: boolean) -> boolean': + entries: + false or false: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + false or null: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + false or true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null or false: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + null or null: + actualResult: false + equivalence: "EQUAL" + error: null + expectedResult: false + null or true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + true or false: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + true or null: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + true or true: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null diff --git a/pinot-integration-tests/src/test/resources/udf-test-results/plus.yaml b/pinot-integration-tests/src/test/resources/udf-test-results/plus.yaml new file mode 100644 index 000000000000..4b89f3750cab --- /dev/null +++ b/pinot-integration-tests/src/test/resources/udf-test-results/plus.yaml @@ -0,0 +1,520 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# This file is auto-generated by the UDF test framework. Do not edit it manually. +# Use the org.apache.pinot.integration.tests.udfUdfTest.generateSnapshots() method to regenerate it. + +--- +Ingestion time transformer: + '(arg0: big_decimal, arg1: big_decimal) -> big_decimal': + entries: + "1 + 2_big_decimal": + actualResult: 3.0 + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 3.0 + "1 + null_big_decimal": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null + '(arg0: double, arg1: double) -> double': + entries: + "1 + 2_double": + actualResult: 3.0 + equivalence: "EQUAL" + error: null + expectedResult: 3.0 + "1 + null_double": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null + '(arg0: float, arg1: float) -> float': + entries: + "1 + 2_float": + actualResult: 3.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 3.0 + "1 + null_float": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null + '(arg0: int, arg1: int) -> int': + entries: + "1 + 2_int": + actualResult: 3.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 3 + "1 + null_int": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null + '(arg0: long, arg1: long) -> long': + entries: + "1 + 2_long": + actualResult: 3.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 3 + "1 + null_long": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null +MSE intermediate stage (with null handling): + '(arg0: big_decimal, arg1: big_decimal) -> big_decimal': + entries: + "1 + 2_big_decimal": + actualResult: 3.0 + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 3.0 + "1 + null_big_decimal": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null + '(arg0: double, arg1: double) -> double': + entries: + "1 + 2_double": + actualResult: 3.0 + equivalence: "EQUAL" + error: null + expectedResult: 3.0 + "1 + null_double": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null + '(arg0: float, arg1: float) -> float': + entries: + "1 + 2_float": + actualResult: 3.0 + equivalence: "EQUAL" + error: null + expectedResult: 3.0 + "1 + null_float": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null + '(arg0: int, arg1: int) -> int': + entries: + "1 + 2_int": + actualResult: 3 + equivalence: "EQUAL" + error: null + expectedResult: 3 + "1 + null_int": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null + '(arg0: long, arg1: long) -> long': + entries: + "1 + 2_long": + actualResult: 3 + equivalence: "EQUAL" + error: null + expectedResult: 3 + "1 + null_long": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null +MSE intermediate stage (without null handling): + '(arg0: big_decimal, arg1: big_decimal) -> big_decimal': + entries: + "1 + 2_big_decimal": + actualResult: 3.0 + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 3.0 + "1 + null_big_decimal": + actualResult: 1.0 + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 1.0 + error: false + errorMessage: null + '(arg0: double, arg1: double) -> double': + entries: + "1 + 2_double": + actualResult: 3.0 + equivalence: "EQUAL" + error: null + expectedResult: 3.0 + "1 + null_double": + actualResult: 1.0 + equivalence: "EQUAL" + error: null + expectedResult: 1.0 + error: false + errorMessage: null + '(arg0: float, arg1: float) -> float': + entries: + "1 + 2_float": + actualResult: 3.0 + equivalence: "EQUAL" + error: null + expectedResult: 3.0 + "1 + null_float": + actualResult: 1.0 + equivalence: "EQUAL" + error: null + expectedResult: 1.0 + error: false + errorMessage: null + '(arg0: int, arg1: int) -> int': + entries: + "1 + 2_int": + actualResult: 3 + equivalence: "EQUAL" + error: null + expectedResult: 3 + "1 + null_int": + actualResult: 1 + equivalence: "EQUAL" + error: null + expectedResult: 1 + error: false + errorMessage: null + '(arg0: long, arg1: long) -> long': + entries: + "1 + 2_long": + actualResult: 3 + equivalence: "EQUAL" + error: null + expectedResult: 3 + "1 + null_long": + actualResult: 1 + equivalence: "EQUAL" + error: null + expectedResult: 1 + error: false + errorMessage: null +SSE predicate (with null handling): + '(arg0: big_decimal, arg1: big_decimal) -> big_decimal': + entries: + "1 + 2_big_decimal": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + "1 + null_big_decimal": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null + '(arg0: double, arg1: double) -> double': + entries: + "1 + 2_double": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + "1 + null_double": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null + '(arg0: float, arg1: float) -> float': + entries: + "1 + 2_float": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + "1 + null_float": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null + '(arg0: int, arg1: int) -> int': + entries: + "1 + 2_int": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + "1 + null_int": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null + '(arg0: long, arg1: long) -> long': + entries: + "1 + 2_long": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + "1 + null_long": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE predicate (without null handling): + '(arg0: big_decimal, arg1: big_decimal) -> big_decimal': + entries: + "1 + 2_big_decimal": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + "1 + null_big_decimal": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null + '(arg0: double, arg1: double) -> double': + entries: + "1 + 2_double": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + "1 + null_double": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null + '(arg0: float, arg1: float) -> float': + entries: + "1 + 2_float": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + "1 + null_float": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null + '(arg0: int, arg1: int) -> int': + entries: + "1 + 2_int": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + "1 + null_int": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null + '(arg0: long, arg1: long) -> long': + entries: + "1 + 2_long": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + "1 + null_long": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE projection (with null handling): + '(arg0: big_decimal, arg1: big_decimal) -> big_decimal': + entries: + "1 + 2_big_decimal": + actualResult: "3.0" + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 3.0 + "1 + null_big_decimal": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null + '(arg0: double, arg1: double) -> double': + entries: + "1 + 2_double": + actualResult: 3.0 + equivalence: "EQUAL" + error: null + expectedResult: 3.0 + "1 + null_double": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null + '(arg0: float, arg1: float) -> float': + entries: + "1 + 2_float": + actualResult: 3.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 3.0 + "1 + null_float": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null + '(arg0: int, arg1: int) -> int': + entries: + "1 + 2_int": + actualResult: 3.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 3 + "1 + null_int": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null + '(arg0: long, arg1: long) -> long': + entries: + "1 + 2_long": + actualResult: 3.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 3 + "1 + null_long": + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null +SSE projection (without null handling): + '(arg0: big_decimal, arg1: big_decimal) -> big_decimal': + entries: + "1 + 2_big_decimal": + actualResult: "3.0" + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 3.0 + "1 + null_big_decimal": + actualResult: "1.0" + equivalence: "BIG_DECIMAL_AS_DOUBLE" + error: null + expectedResult: 1.0 + error: false + errorMessage: null + '(arg0: double, arg1: double) -> double': + entries: + "1 + 2_double": + actualResult: 3.0 + equivalence: "EQUAL" + error: null + expectedResult: 3.0 + "1 + null_double": + actualResult: 1.0 + equivalence: "EQUAL" + error: null + expectedResult: 1.0 + error: false + errorMessage: null + '(arg0: float, arg1: float) -> float': + entries: + "1 + 2_float": + actualResult: 3.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 3.0 + "1 + null_float": + actualResult: 1.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 1.0 + error: false + errorMessage: null + '(arg0: int, arg1: int) -> int': + entries: + "1 + 2_int": + actualResult: 3.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 3 + "1 + null_int": + actualResult: 1.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 1 + error: false + errorMessage: null + '(arg0: long, arg1: long) -> long': + entries: + "1 + 2_long": + actualResult: 3.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 3 + "1 + null_long": + actualResult: 1.0 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 1 + error: false + errorMessage: null diff --git a/pinot-integration-tests/src/test/resources/udf-test-results/toDateTime.yaml b/pinot-integration-tests/src/test/resources/udf-test-results/toDateTime.yaml new file mode 100644 index 000000000000..6e6b7a80a9c2 --- /dev/null +++ b/pinot-integration-tests/src/test/resources/udf-test-results/toDateTime.yaml @@ -0,0 +1,128 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# This file is auto-generated by the UDF test framework. Do not edit it manually. +# Use the org.apache.pinot.integration.tests.udfUdfTest.generateSnapshots() method to regenerate it. + +--- +Ingestion time transformer: + '(mills: long, format: string) -> string': + entries: + Date only: + actualResult: "2020-01-01" + equivalence: "EQUAL" + error: null + expectedResult: "2020-01-01" + UTC ISO8601: + actualResult: "2020-01-01T00:00:00Z" + equivalence: "EQUAL" + error: null + expectedResult: "2020-01-01T00:00:00Z" + error: false + errorMessage: null +MSE intermediate stage (with null handling): + '(mills: long, format: string) -> string': + entries: + Date only: + actualResult: "2020-01-01" + equivalence: "EQUAL" + error: null + expectedResult: "2020-01-01" + UTC ISO8601: + actualResult: "2020-01-01T00:00:00Z" + equivalence: "EQUAL" + error: null + expectedResult: "2020-01-01T00:00:00Z" + error: false + errorMessage: null +MSE intermediate stage (without null handling): + '(mills: long, format: string) -> string': + entries: + Date only: + actualResult: "2020-01-01" + equivalence: "EQUAL" + error: null + expectedResult: "2020-01-01" + UTC ISO8601: + actualResult: "2020-01-01T00:00:00Z" + equivalence: "EQUAL" + error: null + expectedResult: "2020-01-01T00:00:00Z" + error: false + errorMessage: null +SSE predicate (with null handling): + '(mills: long, format: string) -> string': + entries: + Date only: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + UTC ISO8601: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE predicate (without null handling): + '(mills: long, format: string) -> string': + entries: + Date only: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + UTC ISO8601: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE projection (with null handling): + '(mills: long, format: string) -> string': + entries: + Date only: + actualResult: "2020-01-01" + equivalence: "EQUAL" + error: null + expectedResult: "2020-01-01" + UTC ISO8601: + actualResult: "2020-01-01T00:00:00Z" + equivalence: "EQUAL" + error: null + expectedResult: "2020-01-01T00:00:00Z" + error: false + errorMessage: null +SSE projection (without null handling): + '(mills: long, format: string) -> string': + entries: + Date only: + actualResult: "2020-01-01" + equivalence: "EQUAL" + error: null + expectedResult: "2020-01-01" + UTC ISO8601: + actualResult: "2020-01-01T00:00:00Z" + equivalence: "EQUAL" + error: null + expectedResult: "2020-01-01T00:00:00Z" + error: false + errorMessage: null diff --git a/pinot-integration-tests/src/test/resources/udf-test-results/year.yaml b/pinot-integration-tests/src/test/resources/udf-test-results/year.yaml new file mode 100644 index 000000000000..33ad4b321eda --- /dev/null +++ b/pinot-integration-tests/src/test/resources/udf-test-results/year.yaml @@ -0,0 +1,163 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# This file is auto-generated by the UDF test framework. Do not edit it manually. +# Use the org.apache.pinot.integration.tests.udfUdfTest.generateSnapshots() method to regenerate it. + +--- +Ingestion time transformer: + '(millis: long) -> int': + entries: + "1970-01-01T00:00:00Z": + actualResult: 1970 + equivalence: "EQUAL" + error: null + expectedResult: 1970 + "2020-01-01T00:00:00Z": + actualResult: 2020 + equivalence: "EQUAL" + error: null + expectedResult: 2020 + null input: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null +MSE intermediate stage (with null handling): + '(millis: long) -> int': + entries: + "1970-01-01T00:00:00Z": + actualResult: 1970 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 1970 + "2020-01-01T00:00:00Z": + actualResult: 2020 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 2020 + null input: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null +MSE intermediate stage (without null handling): + '(millis: long) -> int': + entries: + "1970-01-01T00:00:00Z": + actualResult: 1970 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 1970 + "2020-01-01T00:00:00Z": + actualResult: 2020 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 2020 + null input: + actualResult: 1970 + equivalence: "NUMBER_AS_DOUBLE" + error: null + expectedResult: 1970 + error: false + errorMessage: null +SSE predicate (with null handling): + '(millis: long) -> int': + entries: + "1970-01-01T00:00:00Z": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + "2020-01-01T00:00:00Z": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null input: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE predicate (without null handling): + '(millis: long) -> int': + entries: + "1970-01-01T00:00:00Z": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + "2020-01-01T00:00:00Z": + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + null input: + actualResult: true + equivalence: "EQUAL" + error: null + expectedResult: true + error: false + errorMessage: null +SSE projection (with null handling): + '(millis: long) -> int': + entries: + "1970-01-01T00:00:00Z": + actualResult: 1970 + equivalence: "EQUAL" + error: null + expectedResult: 1970 + "2020-01-01T00:00:00Z": + actualResult: 2020 + equivalence: "EQUAL" + error: null + expectedResult: 2020 + null input: + actualResult: null + equivalence: "EQUAL" + error: null + expectedResult: null + error: false + errorMessage: null +SSE projection (without null handling): + '(millis: long) -> int': + entries: + "1970-01-01T00:00:00Z": + actualResult: 1970 + equivalence: "EQUAL" + error: null + expectedResult: 1970 + "2020-01-01T00:00:00Z": + actualResult: 2020 + equivalence: "EQUAL" + error: null + expectedResult: 2020 + null input: + actualResult: 1970 + equivalence: "EQUAL" + error: null + expectedResult: 1970 + error: false + errorMessage: null diff --git a/pinot-plugins/pinot-input-format/pinot-avro-base/src/main/java/org/apache/pinot/plugin/inputformat/avro/AvroUtils.java b/pinot-plugins/pinot-input-format/pinot-avro-base/src/main/java/org/apache/pinot/plugin/inputformat/avro/AvroUtils.java index 7fbc087fcb92..4f2c83b9ee22 100644 --- a/pinot-plugins/pinot-input-format/pinot-avro-base/src/main/java/org/apache/pinot/plugin/inputformat/avro/AvroUtils.java +++ b/pinot-plugins/pinot-input-format/pinot-avro-base/src/main/java/org/apache/pinot/plugin/inputformat/avro/AvroUtils.java @@ -27,9 +27,11 @@ import java.util.concurrent.TimeUnit; import java.util.zip.GZIPInputStream; import javax.annotation.Nullable; +import org.apache.avro.Conversions; import org.apache.avro.Schema.Field; import org.apache.avro.SchemaBuilder; import org.apache.avro.file.DataFileStream; +import org.apache.avro.generic.GenericData; import org.apache.avro.generic.GenericDatumReader; import org.apache.avro.generic.GenericRecord; import org.apache.pinot.spi.config.table.ingestion.ComplexTypeConfig; @@ -46,6 +48,21 @@ * Utils for handling Avro records */ public class AvroUtils { + private static final GenericData GENERIC_DATA = new GenericData(); + + static { + // Decimal without scale and precision. Deserialized as BigDecimal, serialized as bytes. + GENERIC_DATA.addLogicalTypeConversion(new Conversions.BigDecimalConversion()); + // Decimal with scale and precision. Deserialized as BigDecimal, serialized as bytes. + GENERIC_DATA.addLogicalTypeConversion(new Conversions.DecimalConversion()); + // TODO: Other interesting standard conversions we may want to add. First we need to make sure that we support + // the corresponding data types in Pinot (ie can we read UUIDs or Instant?). + // UUID is deserialized as java.util.UUID, serialized as string. + //GENERIC_DATA.addLogicalTypeConversion(new Conversions.UUIDConversion()); + // Instant is deserialized as java.time.Instant, serialized as long (epoch millis). + //GENERIC_DATA.addLogicalTypeConversion(new TimeConversions.TimestampMillisConversion()); + } + private AvroUtils() { } @@ -210,15 +227,20 @@ public static org.apache.avro.Schema getAvroSchemaFromPinotSchema(Schema pinotSc return fieldAssembler.endRecord(); } + public static GenericData getGenericData() { + return GENERIC_DATA; + } + /** * Get the Avro file reader for the given file. */ public static DataFileStream getAvroReader(File avroFile) throws IOException { + GenericDatumReader datumReader = new GenericDatumReader<>(null, null, GENERIC_DATA); if (RecordReaderUtils.isGZippedFile(avroFile)) { - return new DataFileStream<>(new GZIPInputStream(new FileInputStream(avroFile)), new GenericDatumReader<>()); + return new DataFileStream<>(new GZIPInputStream(new FileInputStream(avroFile)), datumReader); } else { - return new DataFileStream<>(new FileInputStream(avroFile), new GenericDatumReader<>()); + return new DataFileStream<>(new FileInputStream(avroFile), datumReader); } } diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/AbsUdf.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/AbsUdf.java new file mode 100644 index 000000000000..ddf09a68c8d9 --- /dev/null +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/AbsUdf.java @@ -0,0 +1,63 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.runtime.function; + +import com.google.auto.service.AutoService; +import java.util.Map; +import java.util.Set; +import org.apache.commons.lang3.tuple.Pair; +import org.apache.pinot.common.function.TransformFunctionType; +import org.apache.pinot.common.function.scalar.ArithmeticFunctions; +import org.apache.pinot.core.operator.transform.function.SingleParamMathTransformFunction; +import org.apache.pinot.core.operator.transform.function.TransformFunction; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.core.udf.UdfExampleBuilder; +import org.apache.pinot.core.udf.UdfSignature; + + +@AutoService(Udf.class) +public class AbsUdf extends Udf.FromAnnotatedMethod { + + public AbsUdf() + throws NoSuchMethodException { + super(ArithmeticFunctions.class.getMethod("abs", double.class)); + } + + @Override + public String getDescription() { + return "Returns the absolute value of a numeric input."; + } + + @Override + public Map> getExamples() { + return UdfExampleBuilder.forEndomorphismNumeric(1) + .addExample("positive value", 5, 5) + .addExample("negative value", -3, 3) + .addExample("zero", 0, 0) + .addExample(UdfExample.create("null input", null, null).withoutNull(0)) + .build() + .generateExamples(); + } + + @Override + public Pair> getTransformFunction() { + return Pair.of(TransformFunctionType.ABS, SingleParamMathTransformFunction.AbsTransformFunction.class); + } +} diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/AcosUdf.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/AcosUdf.java new file mode 100644 index 000000000000..c2d1d742ec92 --- /dev/null +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/AcosUdf.java @@ -0,0 +1,73 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.runtime.function; + +import com.google.auto.service.AutoService; +import java.util.Map; +import java.util.Set; +import org.apache.commons.lang3.tuple.Pair; +import org.apache.pinot.common.function.TransformFunctionType; +import org.apache.pinot.common.function.scalar.TrigonometricFunctions; +import org.apache.pinot.core.operator.transform.function.TransformFunction; +import org.apache.pinot.core.operator.transform.function.TrigonometricTransformFunctions.AcosTransformFunction; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.core.udf.UdfExampleBuilder; +import org.apache.pinot.core.udf.UdfParameter; +import org.apache.pinot.core.udf.UdfSignature; +import org.apache.pinot.spi.data.FieldSpec; + + +@AutoService(Udf.class) +public class AcosUdf extends Udf.FromAnnotatedMethod { + + public AcosUdf() + throws NoSuchMethodException { + super(TrigonometricFunctions.class.getMethod("acos", double.class)); + } + + @Override + public String getDescription() { + return "Returns the arc cosine of a numeric input (in radians)."; + } + + @Override + public Map> getExamples() { + return UdfExampleBuilder.forSignature(UdfSignature.of( + UdfParameter.of("d", FieldSpec.DataType.DOUBLE) + .withDescription("The double value for which to compute the arc cosine."), + UdfParameter.result(FieldSpec.DataType.DOUBLE) + .withDescription("The arc cosine of the input value, in radians.") + )) + .addExample("acos(1)", 1d, 0d) + .addExample("acos(0)", 0d, Math.PI / 2) + .addExample("acos(-1)", -1d, Math.PI) + .addExample("larger than 1", 10d, Double.NaN) + .addExample("smaller than -1", -10d, Double.NaN) + .addExample(UdfExample.create("null input", null, null) + .withoutNull(Math.PI / 2)) + .build() + .generateExamples(); + } + + @Override + public Pair> getTransformFunction() { + return Pair.of(TransformFunctionType.ACOS, AcosTransformFunction.class); + } +} diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/Adler32Udf.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/Adler32Udf.java new file mode 100644 index 000000000000..555c0c3b921b --- /dev/null +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/Adler32Udf.java @@ -0,0 +1,66 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.runtime.function; + +import com.google.auto.service.AutoService; +import java.util.Map; +import java.util.Set; +import org.apache.pinot.common.function.scalar.HashFunctions; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.core.udf.UdfExampleBuilder; +import org.apache.pinot.core.udf.UdfParameter; +import org.apache.pinot.core.udf.UdfSignature; +import org.apache.pinot.spi.data.FieldSpec; + + +/** + * UDF stub for adler32 (not implemented). + */ +@AutoService(Udf.class) +public class Adler32Udf extends Udf.FromAnnotatedMethod { + + public Adler32Udf() + throws NoSuchMethodException { + super(HashFunctions.class.getMethod("adler32", byte[].class)); + } + + @Override + public String getDescription() { + return "Computes the Adler-32 checksum of a byte array. " + + "Adler-32 is a checksum algorithm that is simple and fast, but not cryptographically secure."; + } + + @Override + public Map> getExamples() { + return UdfExampleBuilder.forSignature( + UdfSignature.of( + UdfParameter.of("input", FieldSpec.DataType.BYTES) + .withDescription("Input byte array to compute the Adler-32 checksum"), + UdfParameter.result(FieldSpec.DataType.INT) + .withDescription("The Adler-32 checksum of the input byte array") + )) + .addExample("empty", new byte[0], 1) + .addExample("single byte", new byte[]{1}, 131074) + .addExample("multiple bytes", new byte[]{1, 2, 3, 4}, 1572875) + .addExample(UdfExample.create("null input", null, null).withoutNull(1)) + .build() + .generateExamples(); + } +} diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/AndUdf.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/AndUdf.java new file mode 100644 index 000000000000..3218bda441a7 --- /dev/null +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/AndUdf.java @@ -0,0 +1,82 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.runtime.function; + +import com.google.auto.service.AutoService; +import java.util.Map; +import java.util.Set; +import org.apache.commons.lang3.tuple.Pair; +import org.apache.pinot.common.function.PinotScalarFunction; +import org.apache.pinot.common.function.TransformFunctionType; +import org.apache.pinot.core.operator.transform.function.AndOperatorTransformFunction; +import org.apache.pinot.core.operator.transform.function.TransformFunction; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.core.udf.UdfExampleBuilder; +import org.apache.pinot.core.udf.UdfParameter; +import org.apache.pinot.core.udf.UdfSignature; +import org.apache.pinot.spi.data.FieldSpec; + + +/** + * UDF stub for and (not implemented). + */ +@AutoService(Udf.class) +public class AndUdf extends Udf { + @Override + public String getMainName() { + return "and"; + } + + @Override + public String getDescription() { + return "Stub for and function. Not implemented."; + } + + @Override + public Map> getExamples() { + return UdfExampleBuilder.forSignature(UdfSignature.of( + UdfParameter.of("left", FieldSpec.DataType.BOOLEAN) + .withDescription("Left operand of the AND operation"), + UdfParameter.of("right", FieldSpec.DataType.BOOLEAN) + .withDescription("Right operand of the AND operation"), + UdfParameter.result(FieldSpec.DataType.BOOLEAN) + .withDescription("Result of the AND operation, true if both operands are true, false otherwise") + )) + .addExample("true and true", true, true, true) + .addExample("true and false", true, false, false) + .addExample("false and true", false, true, false) + .addExample("false and false", false, false, false) + .addExample(UdfExample.create("true and null", true, null, null).withoutNull(false)) + .addExample(UdfExample.create("null and true", null, true, null).withoutNull(false)) + .addExample(UdfExample.create("null and null", null, null, null).withoutNull(false)) + .build() + .generateExamples(); + } + + @Override + public PinotScalarFunction getScalarFunction() { + return null; + } + + @Override + public Pair> getTransformFunction() { + return Pair.of(TransformFunctionType.AND, AndOperatorTransformFunction.class); + } +} diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/ArrayAverageUdf.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/ArrayAverageUdf.java new file mode 100644 index 000000000000..8f95738f586f --- /dev/null +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/ArrayAverageUdf.java @@ -0,0 +1,76 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.runtime.function; + +import com.google.auto.service.AutoService; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.commons.lang3.tuple.Pair; +import org.apache.pinot.common.function.PinotScalarFunction; +import org.apache.pinot.common.function.TransformFunctionType; +import org.apache.pinot.core.operator.transform.function.ArrayAverageTransformFunction; +import org.apache.pinot.core.operator.transform.function.TransformFunction; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.core.udf.UdfExampleBuilder; +import org.apache.pinot.core.udf.UdfParameter; +import org.apache.pinot.core.udf.UdfSignature; +import org.apache.pinot.spi.data.FieldSpec; + + +@AutoService(Udf.class) +public class ArrayAverageUdf extends Udf { + @Override + public String getMainName() { + return "arrayaverage"; + } + + @Override + public String getDescription() { + return "Returns the average of the elements in the input array."; + } + + @Override + public Map> getExamples() { + return UdfExampleBuilder.forSignature(UdfSignature.of( + UdfParameter.of("input", FieldSpec.DataType.DOUBLE) + .asMultiValued() + .withDescription("Input array of doubles"), + UdfParameter.result(FieldSpec.DataType.DOUBLE) + .withDescription("Average of the input array") + )) + .addExample("average of [1, 2, 3, 4]", List.of(1d, 2d, 3d, 4d), 2.5) + .addExample("average of [5]", List.of(5d), 5d) + .addExample("empty array", List.of(), null) + .addExample(UdfExample.create("null input", null, null).withoutNull(Double.NEGATIVE_INFINITY)) + .build() + .generateExamples(); + } + + @Override + public PinotScalarFunction getScalarFunction() { + return null; + } + + @Override + public Pair> getTransformFunction() { + return Pair.of(TransformFunctionType.ARRAY_AVERAGE, ArrayAverageTransformFunction.class); + } +} diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/ArrayConcatDoubleUdf.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/ArrayConcatDoubleUdf.java new file mode 100644 index 000000000000..6cbf9ebf21ae --- /dev/null +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/ArrayConcatDoubleUdf.java @@ -0,0 +1,71 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.runtime.function; + +import com.google.auto.service.AutoService; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.pinot.common.function.scalar.ArrayFunctions; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.core.udf.UdfExampleBuilder; +import org.apache.pinot.core.udf.UdfParameter; +import org.apache.pinot.core.udf.UdfSignature; +import org.apache.pinot.spi.data.FieldSpec; + + +@AutoService(Udf.class) +public class ArrayConcatDoubleUdf extends Udf.FromAnnotatedMethod { + public ArrayConcatDoubleUdf() + throws NoSuchMethodException { + super(ArrayFunctions.class.getMethod("arrayConcatDouble", double[].class, double[].class)); + } + + @Override + public String getDescription() { + return "Concatenates two arrays of doubles into a single array."; + } + + @Override + public Map> getExamples() { + return UdfExampleBuilder.forSignature(UdfSignature.of( + UdfParameter.of("value1", FieldSpec.DataType.DOUBLE) + .asMultiValued() + .withDescription("First array to concatenate"), + UdfParameter.of("value2", FieldSpec.DataType.DOUBLE) + .asMultiValued() + .withDescription("Second array to concatenate"), + UdfParameter.result(FieldSpec.DataType.DOUBLE) + .asMultiValued() + .withDescription("Concatenated array") + )) + .addExample("two not empty arrays", List.of(1d, 3d), List.of(2d, 4d), List.of(1d, 3d, 2d, 4d)) + .addExample("empty with not empty", List.of(), List.of(1d, 2d), List.of(1d, 2d)) + .addExample("not empty with empty", List.of(1d, 2d), List.of(), List.of(1d, 2d)) + .addExample("two empty arrays", List.of(), List.of(), List.of()) + .addExample(UdfExample.create("null concat not null", null, List.of(1d, 2d), null) + .withoutNull(List.of(1d, 2d))) + .addExample(UdfExample.create("not null concat null", List.of(1d, 2d), null, null) + .withoutNull(List.of(1d, 2d))) + .addExample(UdfExample.create("null input", null, null, null).withoutNull(List.of())) + .build() + .generateExamples(); + } +} diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/ArrayElementAtIntUdf.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/ArrayElementAtIntUdf.java new file mode 100644 index 000000000000..b98f32be2688 --- /dev/null +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/ArrayElementAtIntUdf.java @@ -0,0 +1,65 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.runtime.function; + +import com.google.auto.service.AutoService; +import java.util.Map; +import java.util.Set; +import org.apache.pinot.common.function.scalar.ArrayFunctions; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.core.udf.UdfExampleBuilder; +import org.apache.pinot.core.udf.UdfParameter; +import org.apache.pinot.core.udf.UdfSignature; +import org.apache.pinot.spi.data.FieldSpec; + + +@AutoService(Udf.class) +public class ArrayElementAtIntUdf extends Udf.FromAnnotatedMethod { + public ArrayElementAtIntUdf() + throws NoSuchMethodException { + super(ArrayFunctions.class.getMethod("arrayElementAtInt", int[].class, int.class)); + } + + // language=markdown + @Override + public String getDescription() { + return "Returns the element at the specified index in an array of integers. " + + "The index is 1-based, meaning that the first element is at index 1. "; + } + + @Override + public Map> getExamples() { + return UdfExampleBuilder.forSignature(UdfSignature.of( + UdfParameter.of("arr", FieldSpec.DataType.INT) + .withDescription("Array of integers to retrieve the element from") + .asMultiValued(), + UdfParameter.of("idx", FieldSpec.DataType.INT) + .withDescription("1-based index of the element to retrieve."), + UdfParameter.result(FieldSpec.DataType.INT) + )) + .addExample("first element", new Integer[]{10, 20, 30}, 1, 10) + .addExample("second element", new Integer[]{10, 20, 30}, 2, 20) + .addExample("negative element", new Integer[]{10, 20, 30}, -1, 0) // Index < 1 returns 0 + .addExample("zero element", new Integer[]{10, 20, 30}, 0, 0) // Index < 1 returns 0 + .addExample("out of bounds element", new Integer[]{10, 20, 30}, 4, 0) // Index > length returns 0 + .build() + .generateExamples(); + } +} diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/ArrayMaxUdf.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/ArrayMaxUdf.java new file mode 100644 index 000000000000..ece73a4c3482 --- /dev/null +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/ArrayMaxUdf.java @@ -0,0 +1,77 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.runtime.function; + +import com.google.auto.service.AutoService; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.commons.lang3.tuple.Pair; +import org.apache.pinot.common.function.PinotScalarFunction; +import org.apache.pinot.common.function.TransformFunctionType; +import org.apache.pinot.core.operator.transform.function.ArrayMaxTransformFunction; +import org.apache.pinot.core.operator.transform.function.TransformFunction; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.core.udf.UdfExampleBuilder; +import org.apache.pinot.core.udf.UdfParameter; +import org.apache.pinot.core.udf.UdfSignature; +import org.apache.pinot.spi.data.FieldSpec; + + +@AutoService(Udf.class) +public class ArrayMaxUdf extends Udf { + @Override + public String getMainName() { + return "arrayMax"; + } + + @Override + public String getDescription() { + return "Given an array with numeric values, this function returns the maximum value in the array. " + + "* asdf "; + } + + @Override + public Map> getExamples() { + return UdfExampleBuilder.forSignature(UdfSignature.of( + UdfParameter.of("arr", FieldSpec.DataType.INT) + .withDescription("Input array of integers") + .asMultiValued(), // Input is an array of INT + UdfParameter.result(FieldSpec.DataType.INT) // Return type is single value INT + )) + .addExample("normal array", List.of(1, 2, 3), 3) + .addExample("negative values", List.of(-5, -2, 0), 0) + .addExample("single element", List.of(42), 42) + .addExample(UdfExample.create("empty array", List.of(), null).withoutNull(-2147483648)) + .addExample(UdfExample.create("null array", null, null).withoutNull(-2147483648)) + .build() + .generateExamples(); + } + + @Override + public Pair> getTransformFunction() { + return Pair.of(TransformFunctionType.ARRAY_MAX, ArrayMaxTransformFunction.class); + } + + @Override + public PinotScalarFunction getScalarFunction() { + return null; + } +} diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/ArrayUdf.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/ArrayUdf.java new file mode 100644 index 000000000000..e92c677875d3 --- /dev/null +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/ArrayUdf.java @@ -0,0 +1,51 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.runtime.function; + +import com.google.auto.service.AutoService; +import java.util.Collections; +import java.util.Map; +import java.util.Set; +import org.apache.pinot.common.function.scalar.ArrayFunctions; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.core.udf.UdfSignature; + + +/** + * UDF for array function. + */ +@AutoService(Udf.class) +public class ArrayUdf extends Udf.FromAnnotatedMethod { + public ArrayUdf() + throws NoSuchMethodException { + super(ArrayFunctions.class.getMethod("arrayValueConstructor", Object[].class)); + } + + @Override + public String getDescription() { + return "Constructs an array from the given arguments. Usage: array(element1, element2, ...)."; + } + + @Override + public Map> getExamples() { + // TODO: Support variadic parameters in UdfExampleBuilder and scenarios + return Collections.emptyMap(); + } +} diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/DegreesUdf.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/DegreesUdf.java new file mode 100644 index 000000000000..06afa9e8e05c --- /dev/null +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/DegreesUdf.java @@ -0,0 +1,72 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.runtime.function; + +import com.google.auto.service.AutoService; +import java.util.Map; +import java.util.Set; +import org.apache.commons.lang3.tuple.Pair; +import org.apache.pinot.common.function.TransformFunctionType; +import org.apache.pinot.common.function.scalar.TrigonometricFunctions; +import org.apache.pinot.core.operator.transform.function.TransformFunction; +import org.apache.pinot.core.operator.transform.function.TrigonometricTransformFunctions.DegreesTransformFunction; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.core.udf.UdfExampleBuilder; +import org.apache.pinot.core.udf.UdfParameter; +import org.apache.pinot.core.udf.UdfSignature; +import org.apache.pinot.spi.data.FieldSpec; + + +@AutoService(Udf.class) +public class DegreesUdf extends Udf.FromAnnotatedMethod { + + public DegreesUdf() + throws NoSuchMethodException { + super(TrigonometricFunctions.class.getMethod("degrees", double.class)); + } + + @Override + public String getDescription() { + return "Converts an angle measured in radians to an approximately equivalent angle measured in degrees."; + } + + @Override + public Map> getExamples() { + return UdfExampleBuilder.forSignature(UdfSignature.of( + UdfParameter.of("radians", FieldSpec.DataType.DOUBLE) + .withDescription("The angle in radians to convert to degrees."), + UdfParameter.result(FieldSpec.DataType.DOUBLE) + .withDescription("The angle in degrees.") + )) + .addExample("degrees(0)", 0d, 0d) + .addExample("degrees(PI/2)", Math.PI / 2, 90d) + .addExample("degrees(PI)", Math.PI, 180d) + .addExample("degrees(2*PI)", 2 * Math.PI, 360d) + .addExample(UdfExample.create("null input", null, null) + .withoutNull(0d)) + .build() + .generateExamples(); + } + + @Override + public Pair> getTransformFunction() { + return Pair.of(TransformFunctionType.DEGREES, DegreesTransformFunction.class); + } +} diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/IsFalseUdf.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/IsFalseUdf.java new file mode 100644 index 000000000000..614a941f6950 --- /dev/null +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/IsFalseUdf.java @@ -0,0 +1,75 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.runtime.function; + +import com.google.auto.service.AutoService; +import java.util.Map; +import java.util.Set; +import javax.annotation.Nullable; +import org.apache.commons.lang3.tuple.Pair; +import org.apache.pinot.common.function.PinotScalarFunction; +import org.apache.pinot.common.function.TransformFunctionType; +import org.apache.pinot.core.operator.transform.function.IsFalseTransformFunction; +import org.apache.pinot.core.operator.transform.function.TransformFunction; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.core.udf.UdfExampleBuilder; +import org.apache.pinot.core.udf.UdfParameter; +import org.apache.pinot.core.udf.UdfSignature; +import org.apache.pinot.spi.data.FieldSpec; + + +@AutoService(Udf.class) +public class IsFalseUdf extends Udf { + @Override + public String getMainName() { + return "isFalse"; + } + + @Override + public String getDescription() { + return "Returns true if the input is false (0), otherwise false. Null inputs return false."; + } + + @Override + public Map> getExamples() { + return UdfExampleBuilder.forSignature(UdfSignature.of( + UdfParameter.of("input", FieldSpec.DataType.INT) + .withDescription("Input value to check if it is false (0)"), + UdfParameter.result(FieldSpec.DataType.BOOLEAN) + )) + .addExample("input is 0 (false)", 0, true) + .addExample("input is 1 (true)", 1, false) + .addExample("input is -1 (not false)", -1, false) + .addExample(UdfExample.create("null input", null, false).withoutNull(true)) + .build() + .generateExamples(); + } + + @Override + public Pair> getTransformFunction() { + return Pair.of(TransformFunctionType.IS_FALSE, IsFalseTransformFunction.class); + } + + @Override + @Nullable + public PinotScalarFunction getScalarFunction() { + return null; + } +} diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/MultUdf.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/MultUdf.java new file mode 100644 index 000000000000..096c7f7563f3 --- /dev/null +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/MultUdf.java @@ -0,0 +1,81 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.runtime.function; + +import com.google.auto.service.AutoService; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.commons.lang3.tuple.Pair; +import org.apache.pinot.common.function.PinotScalarFunction; +import org.apache.pinot.common.function.TransformFunctionType; +import org.apache.pinot.common.function.scalar.arithmetic.MultScalarFunction; +import org.apache.pinot.core.operator.transform.function.MultiplicationTransformFunction; +import org.apache.pinot.core.operator.transform.function.TransformFunction; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.core.udf.UdfExampleBuilder; +import org.apache.pinot.core.udf.UdfSignature; + + +@AutoService(Udf.class) +public class MultUdf extends Udf { + @Override + public String getMainName() { + return "mult"; + } + + @Override + public Set getAllNames() { + return Set.of(getMainName(), "times"); + } + + @Override + public String getDescription() { + return "This function multiplies two numeric values together."; + } + + @Override + public String asSqlCall(String name, List sqlArgValues) { + if (name.equals(getMainCanonicalName())) { + return "(" + String.join(" * ", sqlArgValues) + ")"; + } else { + return super.asSqlCall(name, sqlArgValues); + } + } + + @Override + public Map> getExamples() { + return UdfExampleBuilder.forEndomorphismNumeric(2) + .addExample("2 * 3", 2, 3, 6) + .addExample(UdfExample.create("2 * null", 2, null, null).withoutNull(0)) + .build() + .generateExamples(); + } + + @Override + public Pair> getTransformFunction() { + return Pair.of(TransformFunctionType.MULT, MultiplicationTransformFunction.class); + } + + @Override + public PinotScalarFunction getScalarFunction() { + return new MultScalarFunction(); + } +} diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/NotUdf.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/NotUdf.java new file mode 100644 index 000000000000..754d910e0930 --- /dev/null +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/NotUdf.java @@ -0,0 +1,82 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.runtime.function; + +import com.google.auto.service.AutoService; +import java.util.Map; +import java.util.Set; +import org.apache.commons.lang3.tuple.Pair; +import org.apache.pinot.common.function.PinotScalarFunction; +import org.apache.pinot.common.function.TransformFunctionType; +import org.apache.pinot.common.function.scalar.LogicalFunctions; +import org.apache.pinot.core.operator.transform.function.NotOperatorTransformFunction; +import org.apache.pinot.core.operator.transform.function.TransformFunction; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.core.udf.UdfExampleBuilder; +import org.apache.pinot.core.udf.UdfParameter; +import org.apache.pinot.core.udf.UdfSignature; +import org.apache.pinot.spi.data.FieldSpec; + + +@AutoService(Udf.class) +public class NotUdf extends Udf.FromAnnotatedMethod { + + public NotUdf() + throws NoSuchMethodException { + super(LogicalFunctions.class.getMethod("not", boolean.class)); + } + + @Override + public String getMainName() { + return "not"; + } + + @Override + public String getDescription() { + return "Logical NOT function for a boolean value. Returns true if the argument is false, false if true, and null " + + "if null."; + } + + @Override + public Map> getExamples() { + return UdfExampleBuilder.forSignature( + UdfSignature.of( + UdfParameter.of("value", FieldSpec.DataType.BOOLEAN) + .withDescription("A boolean value to negate."), + UdfParameter.result(FieldSpec.DataType.BOOLEAN) + .withDescription("Returns the logical negation of the input value.") + )) + .addExample("not true", true, false) + .addExample("not false", false, true) + .addExample(UdfExample.create("not null", null, null).withoutNull(true)) + .build() + .generateExamples(); + } + + @Override + public PinotScalarFunction getScalarFunction() { + return null; + } + + @Override + public Pair> getTransformFunction() { + return Pair.of(TransformFunctionType.NOT, NotOperatorTransformFunction.class); + } +} diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/OrUdf.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/OrUdf.java new file mode 100644 index 000000000000..b766cc78e7c8 --- /dev/null +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/OrUdf.java @@ -0,0 +1,82 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.runtime.function; + +import com.google.auto.service.AutoService; +import java.util.Map; +import java.util.Set; +import org.apache.commons.lang3.tuple.Pair; +import org.apache.pinot.common.function.PinotScalarFunction; +import org.apache.pinot.common.function.TransformFunctionType; +import org.apache.pinot.core.operator.transform.function.OrOperatorTransformFunction; +import org.apache.pinot.core.operator.transform.function.TransformFunction; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.core.udf.UdfExampleBuilder; +import org.apache.pinot.core.udf.UdfParameter; +import org.apache.pinot.core.udf.UdfSignature; +import org.apache.pinot.spi.data.FieldSpec; + + +@AutoService(Udf.class) +public class OrUdf extends Udf { + @Override + public String getMainName() { + return "or"; + } + + @Override + public String getDescription() { + return "Logical OR function for two boolean values. Returns true if either argument is true, " + + "false if both are false, and null if both are null or one is null and the other is false."; + } + + @Override + public Map> getExamples() { + return UdfExampleBuilder.forSignature(UdfSignature.of( + UdfParameter.of("left", FieldSpec.DataType.BOOLEAN) + .withDescription("Left operand of the OR operation"), + UdfParameter.of("right", FieldSpec.DataType.BOOLEAN) + .withDescription("Right operand of the OR operation"), + UdfParameter.result(FieldSpec.DataType.BOOLEAN) + .withDescription("Result of the OR operation, true if either operand is true, false otherwise") + )) + .addExample("true or true", true, true, true) + .addExample("true or false", true, false, true) + .addExample("false or true", false, true, true) + .addExample("false or false", false, false, false) + .addExample(UdfExample.create("true or null", true, null, true).withoutNull(true)) + .addExample(UdfExample.create("null or true", null, true, true).withoutNull(true)) + .addExample(UdfExample.create("false or null", false, null, null).withoutNull(false)) + .addExample(UdfExample.create("null or false", null, false, null).withoutNull(false)) + .addExample(UdfExample.create("null or null", null, null, null).withoutNull(false)) + .build() + .generateExamples(); + } + + @Override + public PinotScalarFunction getScalarFunction() { + return null; + } + + @Override + public Pair> getTransformFunction() { + return Pair.of(TransformFunctionType.OR, OrOperatorTransformFunction.class); + } +} diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/PlusUdf.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/PlusUdf.java new file mode 100644 index 000000000000..3a008741c238 --- /dev/null +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/PlusUdf.java @@ -0,0 +1,82 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.runtime.function; + +import com.google.auto.service.AutoService; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.commons.lang3.tuple.Pair; +import org.apache.pinot.common.function.PinotScalarFunction; +import org.apache.pinot.common.function.TransformFunctionType; +import org.apache.pinot.common.function.scalar.arithmetic.PlusScalarFunction; +import org.apache.pinot.core.operator.transform.function.AdditionTransformFunction; +import org.apache.pinot.core.operator.transform.function.TransformFunction; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.core.udf.UdfExampleBuilder; +import org.apache.pinot.core.udf.UdfSignature; + + +@AutoService(Udf.class) +public class PlusUdf extends Udf { + @Override + public String getMainName() { + return "plus"; + } + + @Override + public Set getAllNames() { + return Set.of(getMainName(), "add"); + } + + @Override + public String getDescription() { + return "This function adds two numeric values together. In order to concatenate two strings, use the `concat` " + + "function instead."; + } + + @Override + public String asSqlCall(String name, List sqlArgValues) { + if (name.equals(getMainCanonicalName())) { + return "(" + String.join(" + ", sqlArgValues) + ")"; + } else { + return super.asSqlCall(name, sqlArgValues); + } + } + + @Override + public Map> getExamples() { + return UdfExampleBuilder.forEndomorphismNumeric(2) + .addExample("1 + 2", 1, 2, 3) + .addExample(UdfExample.create("1 + null", 1, null, null).withoutNull(1)) + .build() + .generateExamples(); + } + + @Override + public Pair> getTransformFunction() { + return Pair.of(TransformFunctionType.ADD, AdditionTransformFunction.class); + } + + @Override + public PinotScalarFunction getScalarFunction() { + return new PlusScalarFunction(); + } +} diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/ToDateTimeUdf.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/ToDateTimeUdf.java new file mode 100644 index 000000000000..f76537b64311 --- /dev/null +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/ToDateTimeUdf.java @@ -0,0 +1,67 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.runtime.function; + +import com.google.auto.service.AutoService; +import java.util.Map; +import java.util.Set; +import org.apache.pinot.common.function.scalar.DateTimeFunctions; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.core.udf.UdfExampleBuilder; +import org.apache.pinot.core.udf.UdfParameter; +import org.apache.pinot.core.udf.UdfSignature; +import org.apache.pinot.spi.data.FieldSpec; + + +@AutoService(Udf.class) +public class ToDateTimeUdf extends Udf.FromAnnotatedMethod { + + public ToDateTimeUdf() + throws NoSuchMethodException { + super(DateTimeFunctions.class.getMethod("toDateTime", long.class, String.class)); + } + + @Override + public String getDescription() { + return "Converts epoch millis to a DateTime string represented by the given pattern. " + + "Optionally, a timezone can be provided."; + } + + @Override + public Map> getExamples() { + return UdfExampleBuilder.forSignature( + UdfSignature.of( + UdfParameter.of("mills", FieldSpec.DataType.LONG) + .withDescription("A long value representing epoch millis, " + + "e.g., 1577836800000L for 2020-01-01T00:00:00Z"), + UdfParameter.of("format", FieldSpec.DataType.STRING) + .withDescription("A string literal representing the date format, " + + "e.g., 'yyyy-MM-dd'T'HH:mm:ss'Z' or 'yyyy-MM-dd'") + .asLiteralOnly(), + UdfParameter.result(FieldSpec.DataType.STRING) // Return type is single value STRING + )) + .addExample("UTC ISO8601", 1577836800000L, "yyyy-MM-dd'T'HH:mm:ss'Z'", "2020-01-01T00:00:00Z") + .addExample("Date only", 1577836800000L, "yyyy-MM-dd", "2020-01-01") + //.addExample(UdfExample.create("null millis", null, "yyyy-MM-dd", null).withoutNull("1970-01-01")) + //.addExample("null format", 1577836800000L, null, null) + .build() + .generateExamples(); + } +} diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/YearUdf.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/YearUdf.java new file mode 100644 index 000000000000..15524c61e589 --- /dev/null +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/function/YearUdf.java @@ -0,0 +1,71 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.query.runtime.function; + +import com.google.auto.service.AutoService; +import java.util.Map; +import java.util.Set; +import org.apache.commons.lang3.tuple.Pair; +import org.apache.pinot.common.function.TransformFunctionType; +import org.apache.pinot.common.function.scalar.DateTimeFunctions; +import org.apache.pinot.core.operator.transform.function.DateTimeTransformFunction; +import org.apache.pinot.core.operator.transform.function.TransformFunction; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.core.udf.UdfExampleBuilder; +import org.apache.pinot.core.udf.UdfParameter; +import org.apache.pinot.core.udf.UdfSignature; +import org.apache.pinot.spi.data.FieldSpec; + + +@AutoService(Udf.class) +public class YearUdf extends Udf.FromAnnotatedMethod { + + public YearUdf() + throws NoSuchMethodException { + super(DateTimeFunctions.class.getMethod("year", long.class)); + } + + @Override + public String getDescription() { + return "Returns the year from the given epoch millis in UTC timezone."; + } + + @Override + public Map> getExamples() { + return UdfExampleBuilder.forSignature( + UdfSignature.of( + UdfParameter.of("millis", FieldSpec.DataType.LONG) + .withDescription("A long value representing epoch millis" + + "e.g., 1577836800000L for 2020-01-01T00:00:00Z"), + UdfParameter.result(FieldSpec.DataType.INT) + .withDescription("Returns the year as an integer") + )) + .addExample("2020-01-01T00:00:00Z", 1577836800000L, 2020) + .addExample("1970-01-01T00:00:00Z", 0L, 1970) + .addExample(UdfExample.create("null input", null, null).withoutNull(1970)) + .build() + .generateExamples(); + } + + @Override + public Pair> getTransformFunction() { + return Pair.of(TransformFunctionType.YEAR, DateTimeTransformFunction.Year.class); + } +} diff --git a/pinot-udf-test/pom.xml b/pinot-udf-test/pom.xml new file mode 100644 index 000000000000..9a3d2c14eaf1 --- /dev/null +++ b/pinot-udf-test/pom.xml @@ -0,0 +1,48 @@ + + + + 4.0.0 + + pinot + org.apache.pinot + 1.4.0-SNAPSHOT + + pinot-udf-test + Pinot UDF test + https://pinot.apache.org/ + + ${basedir}/.. + + + + + org.apache.pinot + pinot-core + + + + org.testng + testng + test + + + diff --git a/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/PinotFunctionEnvGenerator.java b/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/PinotFunctionEnvGenerator.java new file mode 100644 index 000000000000..2fd0c1291849 --- /dev/null +++ b/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/PinotFunctionEnvGenerator.java @@ -0,0 +1,369 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.udf.test; + +import com.fasterxml.jackson.databind.JsonNode; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; +import org.apache.arrow.util.Preconditions; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.core.udf.UdfExample.NullHandling; +import org.apache.pinot.core.udf.UdfParameter; +import org.apache.pinot.core.udf.UdfSignature; +import org.apache.pinot.segment.spi.index.StandardIndexes; +import org.apache.pinot.spi.config.table.FieldConfig; +import org.apache.pinot.spi.config.table.IndexConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.spi.utils.BytesUtils; +import org.apache.pinot.spi.utils.JsonUtils; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; + + +/// Class used to generate the [TableConfig] and [Schema] for the Pinot function tests. +public class PinotFunctionEnvGenerator { + + private PinotFunctionEnvGenerator() { + } + + public static void prepareEnvironment(UdfTestCluster cluster, Iterable udfs) { + Map> udfByTableName = new HashMap<>(); + for (Udf udf : udfs) { + String tableName = getTableName(udf); + udfByTableName.computeIfAbsent(tableName, k -> new ArrayList<>()).add(udf); + } + + for (Map.Entry> tableEntry : udfByTableName.entrySet()) { + String tableName = tableEntry.getKey(); + List sameTableUdfs = tableEntry.getValue(); + Schema schema = generateSchema(sameTableUdfs); + cluster.addTable(schema, generateTableConfig(sameTableUdfs)); + + cluster.addRows( + tableName, + schema, + sameTableUdfs.stream() + .flatMap(udf -> udf.getExamples().entrySet().stream() + .flatMap(exampleEntry -> { + UdfSignature signature = exampleEntry.getKey(); + return exampleEntry.getValue().stream() + .map(testCase -> asRow(udf, signature, testCase)); + }) + ) + ); + } + } + + public static String getTableName(Udf udf) { + if (udf.getExamples().keySet().stream() + .flatMap(signature -> signature.getParameters().stream()) + .anyMatch(param -> !param.getConstraints().isEmpty())) { + return "udf_test_" + udf.getMainCanonicalName().replaceAll("[^a-zA-Z0-9]", "_"); + } + return "udf_test"; + } + + public static String getUdfColumnName() { + return "udf"; + } + + public static String getTestColumnName() { + return "test"; + } + + public static String getSignatureColumnName() { + return "signature"; + } + + public static List getArgsForCall(UdfSignature signature) { + List params = signature.getParameters(); + List args = new ArrayList<>(params.size()); + for (int i = 0; i < params.size(); i++) { + args.add(getParameterColumnName(signature, i)); + } + return args; + } + + public static List getArgsForCall(UdfSignature signature, UdfExample example) { + List params = signature.getParameters(); + List args = new ArrayList<>(params.size()); + for (int i = 0; i < params.size(); i++) { + UdfParameter param = params.get(i); + String paramValue; + if (param.isLiteralOnly()) { + Object exampleValue = example.getInputValues().get(i); + if (exampleValue == null) { + paramValue = "NULL"; + } else if (exampleValue.getClass().isArray()) { + paramValue = Arrays.toString((Object[]) exampleValue); + } else if (param.getDataType() == FieldSpec.DataType.STRING) { + paramValue = "'" + exampleValue.toString().replace("'", "''") + "'"; + } else if (param.getDataType() == FieldSpec.DataType.BYTES) { + paramValue = "X('" + BytesUtils.toHexString((byte[]) exampleValue) + "')"; + } else { + paramValue = exampleValue.toString(); + } + } else { + paramValue = getParameterColumnName(signature, i); + } + args.add(paramValue); + } + return args; + } + + public static String getParameterColumnName(UdfSignature signature, int argIndex) { + return getParameterColumnName(signature.getParameters().get(argIndex), argIndex); + } + + public static String getParameterColumnName(UdfParameter param, int argIndex) { + return "arg" + argIndex + + "_" + param.getDataType().name().toLowerCase(Locale.US) + + (param.isMultivalued() ? "_mv" : ""); + } + + public static String getResultColumnName(UdfSignature signature, NullHandling nullHandling) { + return getResultColumnName(signature.getReturnType(), nullHandling); + } + + public static String getResultColumnName(UdfParameter param, NullHandling nullHandling) { + return "result" + + "_" + param.getDataType().name().toLowerCase(Locale.US) + + (param.isMultivalued() ? "_mv" : "") + + (nullHandling == NullHandling.ENABLED ? "_null" : ""); + } + + public static TableConfig generateTableConfig(List udfs) { + Preconditions.checkArgument(!udfs.isEmpty(), "UDFs list cannot be empty"); + TableConfigBuilder tableConfigBuilder = new TableConfigBuilder(TableType.OFFLINE) + .setTableName(getTableName(udfs.get(0))); + Map indexConfigs = new HashMap<>(); + indexConfigs.put(StandardIndexes.inverted().getId(), IndexConfig.ENABLED); + + indexConfigs.put(StandardIndexes.dictionary().getId(), IndexConfig.ENABLED); + JsonNode withDictionary = JsonUtils.objectToJsonNode(indexConfigs); + + tableConfigBuilder.addFieldConfig( + new FieldConfig.Builder(getTestColumnName()) + .withIndexes(withDictionary) + .build() + ); + tableConfigBuilder.addFieldConfig( + new FieldConfig.Builder(getUdfColumnName()) + .withIndexes(withDictionary) + .build() + ); + tableConfigBuilder.addFieldConfig( + new FieldConfig.Builder(getSignatureColumnName()) + .withIndexes(withDictionary) + .build() + ); + Set columnNames = new HashSet<>(); + + for (Udf udf : udfs) { + for (UdfSignature signature : udf.getExamples().keySet()) { + List params = signature.getParameters(); + for (int i = 0; i < params.size(); i++) { + createColumnInTableConfig(params.get(i), i, tableConfigBuilder, columnNames); + } + createResultColsInTableConfig(signature.getReturnType(), tableConfigBuilder, columnNames); + } + } + + return tableConfigBuilder.build(); + } + + private static void createColumnInTableConfig( + UdfParameter parameter, + int argIndex, + TableConfigBuilder tableConfigBuilder, + Set columnNames) { + String columnName = getParameterColumnName(parameter, argIndex); + if (!columnNames.contains(columnName)) { + columnNames.add(columnName); + updateTableConfig(parameter, tableConfigBuilder, columnName); + } + } + + private static void createResultColsInTableConfig( + UdfParameter parameter, + TableConfigBuilder tableConfigBuilder, + Set columnNames) { + createResultColInTableConfig(parameter, tableConfigBuilder, columnNames, NullHandling.DISABLED); + createResultColInTableConfig(parameter, tableConfigBuilder, columnNames, NullHandling.ENABLED); + } + + private static void createResultColInTableConfig( + UdfParameter parameter, + TableConfigBuilder tableConfigBuilder, + Set columnNames, + NullHandling nullHandling) { + String columnName = getResultColumnName(parameter, nullHandling); + if (!columnNames.contains(columnName)) { + columnNames.add(columnName); + updateTableConfig(parameter, tableConfigBuilder, columnName); + } + } + + public static Schema generateSchema(List udfs) { + Preconditions.checkArgument(!udfs.isEmpty(), "UDFs list cannot be empty"); + Schema.SchemaBuilder schemaBuilder = new Schema.SchemaBuilder(); + + schemaBuilder.setSchemaName(getTableName(udfs.get(0))); + schemaBuilder.setEnableColumnBasedNullHandling(true); + schemaBuilder.addDimensionField(getTestColumnName(), FieldSpec.DataType.STRING, field -> { + field.setSingleValueField(true); + field.setNullable(false); + }); + schemaBuilder.addDimensionField(getUdfColumnName(), FieldSpec.DataType.STRING, field -> { + field.setSingleValueField(true); + field.setNullable(false); + }); + schemaBuilder.addDimensionField(getSignatureColumnName(), FieldSpec.DataType.STRING, field -> { + field.setSingleValueField(true); + field.setNullable(false); + }); + + Set columnNames = new HashSet<>(); + for (Udf udf : udfs) { + for (UdfSignature signature : udf.getExamples().keySet()) { + List params = signature.getParameters(); + for (int i = 0; i < params.size(); i++) { + UdfParameter parameter = params.get(i); + createParamColInSchema(parameter, i, schemaBuilder, columnNames); + } + createResultColsInSchema(signature.getReturnType(), schemaBuilder, columnNames); + } + } + return schemaBuilder.build(); + } + + private static void createParamColInSchema(UdfParameter parameter, int argIndex, + Schema.SchemaBuilder schemaBuilder, Set columnNames) { + String columnName = getParameterColumnName(parameter, argIndex); + if (!columnNames.contains(columnName)) { + columnNames.add(columnName); + updateSchema(parameter, schemaBuilder, columnName); + } + } + + private static void createResultColsInSchema(UdfParameter parameter, + Schema.SchemaBuilder schemaBuilder, Set columnNames) { + createResultColInSchema(parameter, schemaBuilder, columnNames, NullHandling.DISABLED); + createResultColInSchema(parameter, schemaBuilder, columnNames, NullHandling.ENABLED); + } + + private static void createResultColInSchema( + UdfParameter parameter, + Schema.SchemaBuilder schemaBuilder, + Set columnNames, + NullHandling nullHandling) { + String columnName = getResultColumnName(parameter, nullHandling); + if (!columnNames.contains(columnName)) { + columnNames.add(columnName); + updateSchema(parameter, schemaBuilder, columnName); + } + } + + public static GenericRow asRow( + Udf udf, + UdfSignature signature, + UdfExample testCase) { + GenericRow row = new GenericRow(); + row.putValue(getUdfColumnName(), udf.getMainCanonicalName()); + row.putValue(getTestColumnName(), testCase.getId()); + row.putValue(getSignatureColumnName(), signature.toString()); + row.putValue(getResultColumnName(signature, NullHandling.DISABLED), testCase.getResult(NullHandling.DISABLED)); + row.putValue(getResultColumnName(signature, NullHandling.ENABLED), testCase.getResult(NullHandling.ENABLED)); + List params = signature.getParameters(); + for (int i = 0; i < params.size(); i++) { + Object argValue = testCase.getInputValues().get(i); + String columnName = getParameterColumnName(signature, i); + if (argValue != null) { + row.putValue(columnName, argValue); + } else { + row.addNullValueField(columnName); + } + } + + return row; + } + + private static void updateSchema(UdfParameter param, Schema.SchemaBuilder schemaBuilder, String columnName) { + if (param.isMultivalued()) { + schemaBuilder.addDimensionField(columnName, param.getDataType(), + fieldSpec -> fieldSpec.setSingleValueField(false)); + } else { + switch (param.getDataType()) { + case INT: + case LONG: + case FLOAT: + case DOUBLE: + case BIG_DECIMAL: + case BYTES: + schemaBuilder.addMetricField(columnName, param.getDataType(), + fieldSpec -> fieldSpec.setSingleValueField(true)); + break; + case BOOLEAN: + case TIMESTAMP: + case STRING: + case JSON: + schemaBuilder.addDimensionField(columnName, param.getDataType(), + fieldSpec -> { + fieldSpec.setSingleValueField(true); + }); + break; + case MAP: + case LIST: + case STRUCT: + case UNKNOWN: + default: + throw new IllegalArgumentException("Unsupported data type: " + param.getDataType()); + } + } + } + + public static void updateTableConfig(UdfParameter param, TableConfigBuilder tableConfigBuilder, String columnName) { + List constraints = param.getConstraints(); + if (!constraints.isEmpty()) { + for (UdfParameter.Constraint constraint : constraints) { + constraint.updateTableConfig(tableConfigBuilder, columnName); + } + return; + } + Map indexConfigs = new HashMap<>(); + indexConfigs.put(StandardIndexes.inverted().getId(), IndexConfig.ENABLED); + JsonNode onlyInverted = JsonUtils.objectToJsonNode(indexConfigs); + + tableConfigBuilder.addFieldConfig( + new FieldConfig.Builder(columnName) + .withIndexes(onlyInverted) + .build() + ); + } +} diff --git a/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/ResultByExample.java b/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/ResultByExample.java new file mode 100644 index 000000000000..bd5249c12687 --- /dev/null +++ b/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/ResultByExample.java @@ -0,0 +1,233 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.udf.test; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.collect.Maps; +import java.util.Map; +import java.util.Objects; +import java.util.SortedMap; +import java.util.TreeMap; +import java.util.function.Function; +import javax.annotation.Nullable; +import org.apache.pinot.core.udf.UdfExample; + +/// The result of executing a bunch of [UdfExample] associated with the same signature. +/// +/// This class encapsulates the results of executing a UDF against multiple examples, providing a structured way to +/// represent both successful and failed executions. +/// +/// There are two sublasses: +/// 1. `Failure`: Represents a failure to execute the UDF, containing an error message. This is used by some scenarios +/// to indicate an error running the scenario itself. This means that the scenario wasn't able to extract the results +/// for each example, and thus the results are not available. +/// 2. `Partial`: Represents an execution of the UDF against multiple examples. For each example, it contains whether +/// the example passed or failed, the expected and actual results and the equivalence between them. +/// +/// Each class has its own [Dto] representation, which can be used to serialize the results into a JSON/YAML format. +/// In order to make this serialization simpler, there is a single `Dto` class that is used to represent both `Partial` +/// and `Failure` results. You can use [#asDto()] to convert a `ResultByExample` instance into its DTO representation +/// and vice versa using [Dto#asResultByExample(Function)]. +public abstract class ResultByExample { + + public abstract Dto asDto(); + + public static class Partial extends ResultByExample { + private final Map _resultsByExample; + private final Map _equivalenceByExample; + private final Map _errorsByExample; + + public Partial(Map resultsByExample, + Map equivalenceByExample, + Map errorsByExample) { + _resultsByExample = resultsByExample; + _equivalenceByExample = equivalenceByExample; + _errorsByExample = errorsByExample; + } + + public Map getResultsByExample() { + return _resultsByExample; + } + + public Map getEquivalenceByExample() { + return _equivalenceByExample; + } + + public Map getErrorsByExample() { + return _errorsByExample; + } + + @Override + public Dto asDto() { + SortedMap dtoEntries = new TreeMap<>(); + for (Map.Entry entry : _resultsByExample.entrySet()) { + UdfExample example = entry.getKey(); + UdfExampleResult result = entry.getValue(); + String error = _errorsByExample.get(example); + dtoEntries.put(example.getId(), + new Dto.DtoEntry(result.getExpectedResult(), result.getActualResult(), + _equivalenceByExample.get(example), error)); + } + return new Dto(dtoEntries, null); + } + + @Override + public boolean equals(Object o) { + if (!(o instanceof Partial)) { + return false; + } + Partial partial = (Partial) o; + return Objects.equals(getResultsByExample(), partial.getResultsByExample()) && Objects.equals( + getEquivalenceByExample(), partial.getEquivalenceByExample()) && Objects.equals(getErrorsByExample(), + partial.getErrorsByExample()); + } + + @Override + public int hashCode() { + return Objects.hash(getResultsByExample(), getEquivalenceByExample(), getErrorsByExample()); + } + } + + public static class Failure extends ResultByExample { + private final String _errorMessage; + + public Failure(String errorMessage) { + _errorMessage = errorMessage; + } + + public String getErrorMessage() { + return _errorMessage; + } + + @Override + public Dto asDto() { + return new Dto(null, _errorMessage); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof Failure)) { + return false; + } + Failure that = (Failure) o; + return _errorMessage.equals(that._errorMessage); + } + + @Override + public int hashCode() { + return _errorMessage.hashCode(); + } + } + + public static class Dto { + @Nullable + private final SortedMap _entries; + @Nullable + private final String _errorMessage; + + @JsonCreator + public Dto( + @Nullable @JsonProperty("entries") SortedMap entries, + @JsonProperty("globalError") @Nullable String errorMessage) { + _entries = entries; + _errorMessage = errorMessage; + } + + @Nullable + public SortedMap getEntries() { + return _entries; + } + + @Nullable + public String getErrorMessage() { + return _errorMessage; + } + + public boolean isError() { + return _errorMessage != null; + } + + public ResultByExample asResultByExample(Function exampleById) { + if (isError()) { + return new Failure(_errorMessage); + } + assert _entries != null; + int size = _entries.size(); + Map resultsByExample = Maps.newHashMapWithExpectedSize(size); + Map equivalenceByExample = Maps.newHashMapWithExpectedSize(size); + Map errorsByExample = Maps.newHashMapWithExpectedSize(size); + + for (Map.Entry entry : _entries.entrySet()) { + String exampleId = entry.getKey(); + DtoEntry dtoEntry = entry.getValue(); + UdfExample example = exampleById.apply(exampleId); + + UdfExampleResult result = dtoEntry.getError() == null + ? UdfExampleResult.success(example, dtoEntry.getActualResult(), dtoEntry.getExpectedResult()) + : UdfExampleResult.error(example, dtoEntry.getError()); + resultsByExample.put(example, result); + equivalenceByExample.put(example, dtoEntry.getEquivalence()); + errorsByExample.put(example, dtoEntry.getError()); + } + + return new Partial(resultsByExample, equivalenceByExample, errorsByExample); + } + + public static class DtoEntry { + private final Object _expectedResult; + private final Object _actualResult; + private final UdfTestFramework.EquivalenceLevel _equivalence; + @Nullable + private final String _error; + + @JsonCreator + public DtoEntry( + @JsonProperty("expectedResult") Object expectedResult, + @JsonProperty("actualResult") Object actualResult, + @JsonProperty("equivalence") UdfTestFramework.EquivalenceLevel equivalence, + @JsonProperty("error") @Nullable String error) { + _expectedResult = expectedResult; + _actualResult = actualResult; + _equivalence = equivalence; + _error = error; + } + + public Object getExpectedResult() { + return _expectedResult; + } + + public Object getActualResult() { + return _actualResult; + } + + public UdfTestFramework.EquivalenceLevel getEquivalence() { + return _equivalence; + } + + @Nullable + public String getError() { + return _error; + } + } + } +} diff --git a/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/UdfExampleResult.java b/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/UdfExampleResult.java new file mode 100644 index 000000000000..1b835c31fd34 --- /dev/null +++ b/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/UdfExampleResult.java @@ -0,0 +1,97 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.udf.test; + +import javax.annotation.Nullable; +import org.apache.pinot.core.udf.UdfExample; + +/// The object used to return results from the UDF test execution framework for a single UDF example. +/// +/// An execution may be successful, in which case error message will be null, or it may fail, +/// in which case the error message will be set. Remember that the actual and expected results may be null even on +/// success. +public class UdfExampleResult { + + private final UdfExample _test; + /// The result of the test execution in case it was successful. + /// + /// Remember this value can be null even on success. + /// + /// The type of the result is determined by [UdfTestScenario] that was performed. While most of the time it will + /// be the returned call value, in predicate executions it may be a boolean indicating whether the value matched the + /// expected result or not + @Nullable + private final Object _actualResult; + /// The expected result of the test execution. + /// + /// Remember this value can be null even on success. + /// + /// The type of the result is determined by [UdfTestScenario] that was performed. While most of the time it will + /// be the returned call value, in predicate executions it may be a boolean indicating whether the value matched the + /// expected result or not + @Nullable + private final Object _expectedResult; + /// The error message in case the test execution failed. + /// If null, the test execution was successful. + @Nullable + private final String _errorMessage; + + private UdfExampleResult( + UdfExample test, + @Nullable Object actualResult, + @Nullable Object expectedResult, + @Nullable String errorMessage) { + _test = test; + _actualResult = actualResult; + _expectedResult = expectedResult; + _errorMessage = errorMessage; + } + + /// Creates a test result indicating a successful test execution. + public static UdfExampleResult success( + UdfExample test, + @Nullable Object actualResult, + @Nullable Object expectedResult) { + return new UdfExampleResult(test, actualResult, expectedResult, null); + } + + /// Creates a test result indicating an error in the test execution. + public static UdfExampleResult error(UdfExample test, String errorMessage) { + return new UdfExampleResult(test, null, null, errorMessage); + } + + public UdfExample getTest() { + return _test; + } + + @Nullable + public Object getActualResult() { + return _actualResult; + } + + @Nullable + public Object getExpectedResult() { + return _expectedResult; + } + + @Nullable + public String getErrorMessage() { + return _errorMessage; + } +} diff --git a/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/UdfReporter.java b/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/UdfReporter.java new file mode 100644 index 000000000000..bf897bb39a86 --- /dev/null +++ b/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/UdfReporter.java @@ -0,0 +1,358 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.udf.test; + +import java.io.PrintWriter; +import java.io.Writer; +import java.util.Arrays; +import java.util.Collection; +import java.util.Comparator; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; +import java.util.SortedMap; +import java.util.TreeMap; +import java.util.TreeSet; +import java.util.function.Function; +import java.util.stream.Collectors; +import javax.annotation.Nullable; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.core.udf.UdfParameter; +import org.apache.pinot.core.udf.UdfSignature; +import org.apache.pinot.spi.utils.BytesUtils; + + +/// A class that generates a markdown report for the UDF test results. +public class UdfReporter { + + private UdfReporter() { + } + + /// Generates a markdown report for the given UDF and its test results. + public static void reportAsMarkdown(Udf udf, UdfTestResult.ByScenario byScenario, Writer writer) { + try (PrintWriter report = new PrintWriter(writer)) { + report.append("## ").append(udf.getMainName()).append("\n\n"); + + if (udf.getAllNames().size() > 1) { + report.append("Other names: ") + .append(udf.getAllNames().stream() + .filter(name -> !name.equals(udf.getMainName())) + .collect(Collectors.joining(", "))) + .append("\n\n"); + } + + report.append("### Description\n\n") + .append(udf.getDescription()).append("\n"); + + TreeSet scenarios = new TreeSet<>(Comparator.comparing(UdfTestScenario::getTitle)); + scenarios.addAll(byScenario.getMap().keySet()); + + reportSummary(udf, byScenario, scenarios, report); + reportSignatures(udf, report); + reportScenarios(udf, byScenario, report, scenarios); + report.flush(); + } + } + + private static void reportScenarios(Udf udf, UdfTestResult.ByScenario byScenario, PrintWriter report, + TreeSet scenarios) { + report.append("### Scenarios\n\n"); + + // This is used to create a collapsed section in the markdown report + report.append("
\n" + + "\n" + + "Click to open\n\n"); + + for (UdfTestScenario scenario : scenarios) { + UdfTestResult.BySignature bySignature = byScenario.getMap().get(scenario); + + TreeSet signatures = new TreeSet<>(Comparator.comparing(UdfSignature::toString)); + signatures.addAll(bySignature.getMap().keySet()); + + report.append("#### ").append(scenario.getTitle()).append("\n\n"); + + report.append('\n'); + + report.append("| Signature | Call | Expected result | Actual result | Comparison or Error |\n"); + report.append("|-----------|------|-----------------|---------------|---------------------|\n"); + + for (UdfSignature signature : signatures) { + ResultByExample resultByExample = bySignature.getMap().get(signature); + + if (resultByExample instanceof ResultByExample.Partial) { + ResultByExample.Partial partial = (ResultByExample.Partial) resultByExample; + for (Map.Entry exampleEntry : partial.getResultsByExample().entrySet()) { + UdfExample example = exampleEntry.getKey(); + UdfExampleResult testResult = exampleEntry.getValue(); + + // Signature column + report.append("| ") + .append(signature.toString()).append(" | "); + + // Call column + report.append(asSqlCallWithLiteralArgs(udf, udf.getMainName(), example.getInputValues())) + .append(" | "); + + // Expected result + Object expected = testResult.getExpectedResult(); + Object actual = testResult.getActualResult(); + + Function valueFormatter = getResultFormatter(expected, actual); + report.append(valueFormatter.apply(expected)).append(" | ") + .append(valueFormatter.apply(actual)).append(" | "); + + // Comparison or Error + String error = partial.getErrorsByExample().get(example); + if (error != null) { + report.append("❌ ").append(error.replace("\n", " ")).append(" |\n"); + } else { + UdfTestFramework.EquivalenceLevel comparison = partial.getEquivalenceByExample().get(example); + report.append(comparison != null ? comparison.name() : "").append(" |\n"); + } + } + } else if (resultByExample instanceof ResultByExample.Failure) { + ResultByExample.Failure failure = (ResultByExample.Failure) resultByExample; + + report.append("| ") + .append(signature.toString()) + .append(" | - | - | - | ❌ ") + .append(failure.getErrorMessage().replace("\n", " ")) + .append(" |\n"); + } + } + report.append("\n"); + } + // Close the collapsed section + report.append("\n
\n\n"); + } + + private static void reportSignatures(Udf udf, PrintWriter report) { + Set signatures = new TreeSet<>(Comparator.comparing(UdfSignature::toString)); + signatures.addAll(udf.getExamples().keySet()); + if (!signatures.isEmpty()) { + report.append("### Signatures\n\n"); + + boolean paramsAlreadyPrinted = false; + for (UdfSignature signature : signatures) { + report.append("#### ").append(udf.getMainName()).append(signature.toString()).append("\n\n"); + + String resultDescription = signature.getReturnType().getDescription(); + if (resultDescription != null) { + report.append(resultDescription).append("\n\n"); + } + + if (signature.getParameters().stream().anyMatch(p -> p.getDescription() != null)) { + // This is used to create a collapsed section in the markdown report + if (paramsAlreadyPrinted) { + report.append("
\n" + + "\n" + + "Click to open\n\n"); + } + + report.append("| Parameter | Type | Description |\n"); + report.append("|-----------|------|-------------|\n"); + for (UdfParameter parameter : signature.getParameters()) { + report.append("| ") + .append(parameter.getName()).append(" | ") + .append(parameter.getDataType().toString().toLowerCase(Locale.US)).append(" | ") + .append(parameter.getDescription() != null ? parameter.getDescription() : "") + .append(" |\n"); + } + if (paramsAlreadyPrinted) { + // Close the collapsed section + report.append("\n
\n\n"); + } else { + paramsAlreadyPrinted = true; + } + } + } + } + } + + private static void reportSummary(Udf udf, UdfTestResult.ByScenario byScenario, TreeSet scenarios, + PrintWriter report) { + SortedMap summaries = new TreeMap<>(Comparator.comparing(UdfTestScenario::getTitle)); + for (UdfTestScenario scenario : scenarios) { + UdfTestResult.BySignature bySignature = byScenario.getMap().get(scenario); + summaries.put(scenario, summarize(bySignature)); + } + report.append("### Summary\n\n"); + + udf.getExamples().keySet().stream() + .min(Comparator.comparing(UdfSignature::toString)) + .ifPresent(udfSignature -> { + + report.append("|Call | Result (with null handling) | Result (without null handling)\n"); + report.append("|-----|-----------------------------|------------------------------|\n"); + + for (UdfExample example : udf.getExamples().get(udfSignature)) { + // Expected result + Object withNull = example.getResult(UdfExample.NullHandling.DISABLED); + Object withoutNull = example.getResult(UdfExample.NullHandling.DISABLED); + + // Call column + report.append("| ") + .append(asSqlCallWithLiteralArgs(udf, udf.getMainName(), example.getInputValues())) + .append(" | ") + .append(valueToString(withNull)) + .append(" | ") + .append(valueToString(withoutNull)) + .append(" |\n"); + } + report.append("\n"); + }); + + if (summaries.values().stream().distinct().count() == 1) { + String summary = summaries.values().iterator().next(); + if (summary.equals("❌ Unsupported")) { + report.append("The UDF ").append(udf.getMainName()) + .append(" is not supported in all scenarios.\n\n"); + } else if (summary.contains("❌")) { + report.append("The UDF ").append(udf.getMainName()) + .append(" has failed in all scenarios with the following error: ") + .append(summary).append("\n\n"); + } else if (summary.equals("EQUAL")) { + report.append("The UDF ").append(udf.getMainName()) + .append(" is supported in all scenarios\n\n"); + } else { + report.append("The UDF ").append(udf.getMainName()) + .append(" is supported in all scenarios with at least ") + .append(summary).append(" semantic.\n\n"); + } + } else { + report.append("This UDF has different semantics in different scenarios:\n\n"); + + report.append("| Scenario | Semantic |\n"); + report.append("|----------|----------|\n"); + + for (Map.Entry entry : summaries.entrySet()) { + report.append("| ").append(entry.getKey().getTitle()).append(" | ") + .append(entry.getValue()).append(" |\n"); + } + } + } + + private static String valueToString(@Nullable Object value) { + if (value == null) { + return "NULL"; + } else if (value.getClass().isArray()) { + if (value.getClass().isAssignableFrom(byte[].class)) { + return "hexToBytes('" + BytesUtils.toHexString((byte[]) value) + "')"; + } + return Arrays.stream((Object[]) value) + .map(UdfReporter::valueToString) + .collect(Collectors.joining(", ", "[", "]")); + } else if (value instanceof String) { + return "'" + value.toString().replace("'", "''") + "'"; + } else if (value instanceof Collection) { + return ((Collection) value).stream() + .map(UdfReporter::valueToString) + .collect(Collectors.joining(", ", "[", "]")); + } else if (value instanceof Number || value instanceof Boolean) { + return value.toString(); + } else { + return value.toString(); + } + } + + private static String summarize(UdfTestResult.BySignature bySignature) { + + UdfTestFramework.EquivalenceLevel comparison = UdfTestFramework.EquivalenceLevel.EQUAL; + boolean withSuccess = false; + int errors = 0; + for (ResultByExample result : bySignature.getMap().values()) { + if (result instanceof ResultByExample.Failure) { + String error = (((ResultByExample.Failure) result)).getErrorMessage().replace("\n", " "); + return "❌ " + error; + } + if (result instanceof ResultByExample.Partial) { + ResultByExample.Partial partial = (ResultByExample.Partial) result; + + withSuccess |= !partial.getEquivalenceByExample().isEmpty(); + for (UdfTestFramework.EquivalenceLevel value : partial.getEquivalenceByExample().values()) { + if (value.compareTo(comparison) > 0) { + comparison = value; + } + } + + errors += partial.getErrorsByExample().values().size(); + } + } + + if (withSuccess) { + if (errors == 0) { + return comparison.name(); + } + return comparison.name() + " with " + errors + " errors."; + } else { + return "Not supported"; + } + } + + private static Function getResultFormatter(Object expected, Object actual) { + Function valueFormatter; + if (expected != null && actual != null && expected.getClass().equals(actual.getClass())) { + valueFormatter = value -> { + if (value.getClass().isArray()) { + return Arrays.toString((Object[]) value); + } + return value.toString(); + }; + } else { + valueFormatter = value -> { + if (value == null) { + return "NULL"; + } else if (value.getClass().isArray()) { + String componentTypeName = value.getClass().getComponentType().getSimpleName(); + String valueDesc; + switch (componentTypeName) { + case "int": + valueDesc = Arrays.toString((int[]) value); + break; + case "long": + valueDesc = Arrays.toString((long[]) value); + break; + case "float": + valueDesc = Arrays.toString((float[]) value); + break; + case "double": + valueDesc = Arrays.toString((double[]) value); + break; + default: + valueDesc = Arrays.toString((Object[]) value); + break; + } + return valueDesc + " ( array of " + componentTypeName + ")"; + } else { + return value + " (" + value.getClass().getSimpleName() + ")"; + } + }; + } + return valueFormatter; + } + + private static String asSqlCallWithLiteralArgs(Udf udf, String name, List inputs) { + List args = inputs.stream() + .map(UdfReporter::valueToString) + .collect(Collectors.toList()); + return udf.asSqlCall(name, args); + } +} diff --git a/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/UdfTestCluster.java b/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/UdfTestCluster.java new file mode 100644 index 000000000000..0d5ae92976c4 --- /dev/null +++ b/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/UdfTestCluster.java @@ -0,0 +1,70 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.udf.test; + +import java.util.Iterator; +import java.util.stream.Stream; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.GenericRow; + + +/// An interface for executing queries in Pinot function tests. +/// +/// For example, one implementation can start a local cluster on the same JVM while another can connect to a remote +/// cluster. +public interface UdfTestCluster extends AutoCloseable { + + void start(); + + /// Adds a table to the cluster with the given schema and table configuration. + /// Implementations can assume that the table doesn not exist in the cluster. + void addTable(Schema schema, TableConfig tableConfig); + + /// Inserts the given rows into the specified table. + void addRows(String tableName, Schema schema, Stream rows); + + /// Executes a query and returns an iterator over the results. + Iterator query(ExecutionContext context, String sql); + + /// Closes the cluster and releases any resources that were allocated on start. + /// Implementations must be sure that any table created in the cluster is removed + @Override + void close() + throws Exception; + + class ExecutionContext { + private final UdfExample.NullHandling _nullHandlingMode; + private final boolean _useMultistageEngine; + + public ExecutionContext(UdfExample.NullHandling nullHandlingMode, boolean useMultistageEngine) { + _nullHandlingMode = nullHandlingMode; + _useMultistageEngine = useMultistageEngine; + } + + public UdfExample.NullHandling getNullHandlingMode() { + return _nullHandlingMode; + } + + public boolean isUseMultistageEngine() { + return _useMultistageEngine; + } + } +} diff --git a/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/UdfTestFramework.java b/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/UdfTestFramework.java new file mode 100644 index 000000000000..421f7c554495 --- /dev/null +++ b/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/UdfTestFramework.java @@ -0,0 +1,288 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.udf.test; + +import com.google.common.collect.Maps; +import java.math.BigDecimal; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Map; +import java.util.Objects; +import java.util.ServiceLoader; +import java.util.Set; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Future; +import java.util.stream.Collectors; +import javax.annotation.Nullable; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.core.udf.UdfSignature; + +/// This is the entry paint for the UDF test framework. It allows to run tests for UDFs using a given cluster. +public class UdfTestFramework { + + private static final EquivalenceLevel[] EQUIVALENCE_LEVELS = new EquivalenceLevel[]{ + EquivalenceLevel.EQUAL, + EquivalenceLevel.BIG_DECIMAL_AS_DOUBLE, + EquivalenceLevel.NUMBER_AS_DOUBLE + }; + private final Set _udfs; + private final UdfTestCluster _cluster; + private final Set _scenarios; + private final ExecutorService _executorService; + + public UdfTestFramework(Set udfs, UdfTestCluster cluster, + Set scenarios, ExecutorService executorService) { + _udfs = udfs; + _cluster = cluster; + _scenarios = scenarios; + _executorService = executorService; + } + + public static UdfTestFramework fromServiceLoader(UdfTestCluster cluster, + ExecutorService executorService) { + Set udfs = ServiceLoader.load(Udf.class).stream() + .map(ServiceLoader.Provider::get) + .collect(Collectors.toSet()); + Set scenarios = ServiceLoader.load(UdfTestScenario.Factory.class).stream() + .map(ServiceLoader.Provider::get) + .map(factory -> factory.create(cluster)) + .collect(Collectors.toSet()); + + return new UdfTestFramework(udfs, cluster, scenarios, executorService); + } + + public Set getUdfs() { + return _udfs; + } + + public Set getScenarios() { + return _scenarios; + } + + public void startUp() { + PinotFunctionEnvGenerator.prepareEnvironment(_cluster, _udfs); + } + + /// Executes all UDFs in all scenarios and returns the results. + public UdfTestResult execute() + throws InterruptedException { + Map>>>> asyncResults + = Maps.newHashMapWithExpectedSize(_udfs.size()); + for (Udf udf : _udfs) { + Map>>> scenarioResults + = Maps.newHashMapWithExpectedSize(_scenarios.size()); + executeAsync(udf, scenarioResults); + asyncResults.put(udf, scenarioResults); + } + + return byUdf(asyncResults); + } + + /// Executes a single UDF in all scenarios and returns the results. + public UdfTestResult.ByScenario execute(Udf udf) + throws InterruptedException { + Map>>> scenarioResults + = Maps.newHashMapWithExpectedSize(_scenarios.size()); + executeAsync(udf, scenarioResults); + + return byScenario(scenarioResults); + } + + private void executeAsync( + Udf udf, + Map>>> scenarioResults + ) { + for (UdfTestScenario scenario : _scenarios) { + Set udfSignatures = udf.getExamples().keySet(); + Map>> signatureResults + = Maps.newHashMapWithExpectedSize(udfSignatures.size()); + for (UdfSignature signature: udfSignatures) { + signatureResults.put(signature, _executorService.submit(() -> scenario.execute(udf, signature))); + } + scenarioResults.put(scenario, signatureResults); + } + } + + private UdfTestResult byUdf( + Map>>>> tasks + ) throws InterruptedException { + Map results = Maps.newHashMapWithExpectedSize(tasks.size()); + for (Map.Entry>>>> entry + : tasks.entrySet()) { + Udf udf = entry.getKey(); + Map>>> scenarioTasks + = entry.getValue(); + results.put(udf, byScenario(scenarioTasks)); + } + return new UdfTestResult(results); + } + + private UdfTestResult.ByScenario byScenario( + Map>>> scenarioTasks + ) throws InterruptedException { + Map results = Maps.newHashMapWithExpectedSize(scenarioTasks.size()); + for (Map.Entry>>> entry + : scenarioTasks.entrySet()) { + UdfTestScenario scenario = entry.getKey(); + Map>> tasks = entry.getValue(); + results.put(scenario, bySignature(tasks)); + } + return new UdfTestResult.ByScenario(results); + } + + private UdfTestResult.BySignature bySignature( + Map>> tasks + ) throws InterruptedException { + Map results = Maps.newHashMapWithExpectedSize(tasks.size()); + for (Map.Entry>> entry : tasks.entrySet()) { + UdfSignature signature = entry.getKey(); + Future> task = entry.getValue(); + ResultByExample result = resolve(task); + results.put(signature, result); + } + return new UdfTestResult.BySignature(results); + } + + private ResultByExample resolve(Future> task) + throws InterruptedException { + try { + Map result = task.get(); + Map comparisons = Maps.newHashMapWithExpectedSize(result.size()); + Map errors = Maps.newHashMapWithExpectedSize(result.size()); + for (Map.Entry entry : result.entrySet()) { + try { + EquivalenceLevel equivalence = compareResult(entry.getValue()); + comparisons.put(entry.getKey(), equivalence); + } catch (Exception e) { + errors.put(entry.getKey(), e.getMessage()); + } + } + return new ResultByExample.Partial(result, comparisons, errors); + } catch (ExecutionException e) { + if (e.getCause().getMessage().contains("Unsupported function")) { + return new ResultByExample.Failure("Unsupported"); + } + return new ResultByExample.Failure(e.getCause().getMessage()); + } + } + + private EquivalenceLevel compareResult(UdfExampleResult result) { + Object actualResult = result.getActualResult(); + Object expectedResult = result.getExpectedResult(); + + for (EquivalenceLevel equivalence : EQUIVALENCE_LEVELS) { + try { + equivalence.check(expectedResult, actualResult); + return equivalence; + } catch (AssertionError e) { + // Continue to the next equivalence if the current one fails + } + } + throw new RuntimeException("Unexpected value"); + } + + public enum EquivalenceLevel { + EQUAL, + BIG_DECIMAL_AS_DOUBLE, + NUMBER_AS_DOUBLE, + ERROR; + + public void check(@Nullable Object expected, @Nullable Object actual) throws AssertionError { + expected = canonizeObject(expected); + actual = canonizeObject(actual); + switch (this) { + case EQUAL: + if (!Objects.equals(expected, actual)) { + throw new AssertionError(describeDiscrepancy(expected, actual)); + } + break; + case NUMBER_AS_DOUBLE: + if (expected == null && actual == null) { + return; // Both are null, considered equal + } + if (expected instanceof BigDecimal && actual instanceof String) { + // Big decimals are sent as strings through the wire protocol, so we need to convert them + try { + actual = new BigDecimal((String) actual); + } catch (NumberFormatException e) { + throw new AssertionError("Expected a number as actual value but got the String: " + actual); + } + } + if (!(expected instanceof Number) || !(actual instanceof Number)) { + Class expectedClass = expected != null ? expected.getClass() : null; + Class actualClass = actual != null ? actual.getClass() : null; + throw new AssertionError("Both expected and actual should be numbers for NUMBER_AS_DOUBLE " + + "comparison, but got: expected=" + expected + "(of type " + expectedClass + ")" + + ", actual=" + actual + "(of type " + actualClass + ")"); + } + if (Double.compare(((Number) expected).doubleValue(), ((Number) actual).doubleValue()) != 0) { + throw new AssertionError(describeDiscrepancy(expected, actual)); + } + break; + case BIG_DECIMAL_AS_DOUBLE: { + if (expected instanceof BigDecimal) { + NUMBER_AS_DOUBLE.check(expected, actual); + } else { + EQUAL.check(expected, actual); + } + break; + } + case ERROR: + throw new UnsupportedOperationException("Comparison type ERROR is not supported"); + default: + throw new IllegalArgumentException("Unknown comparison type: " + this); + } + } + + private static String describeDiscrepancy(@Nullable Object expected, @Nullable Object actual) { + String expectedDesc = expected == null ? "null" : expected + " (" + expected.getClass().getSimpleName() + ")"; + String actualDesc = actual == null ? "null" : actual + " (" + actual.getClass().getSimpleName() + ")"; + return "Expected: " + expectedDesc + ", but got: " + actualDesc; + } + + private static Object canonizeObject(@Nullable Object value) { + if (value == null) { + return null; + } + if (value.getClass().isArray()) { + Class componentType = value.getClass().getComponentType(); + if (componentType == int.class) { + return Arrays.stream((int[]) value).boxed().collect(Collectors.toList()); + } else if (componentType == long.class) { + return Arrays.stream((long[]) value).boxed().collect(Collectors.toList()); + } else if (componentType == double.class) { + return Arrays.stream((double[]) value).boxed().collect(Collectors.toList()); + } else if (componentType == float.class) { + // Convert float array to List for consistency + float[] floatArray = (float[]) value; + ArrayList list = new ArrayList<>(floatArray.length); + for (float f : floatArray) { + list.add(f); + } + return list; + } else { + return Arrays.asList((Object[]) value); + } + } + return value; + } + } +} diff --git a/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/UdfTestResult.java b/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/UdfTestResult.java new file mode 100644 index 000000000000..1bb6af85ab8c --- /dev/null +++ b/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/UdfTestResult.java @@ -0,0 +1,220 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.udf.test; + +import com.fasterxml.jackson.annotation.JsonAnyGetter; +import com.fasterxml.jackson.annotation.JsonAnySetter; +import com.fasterxml.jackson.annotation.JsonCreator; +import com.google.common.collect.Maps; +import java.util.Map; +import java.util.SortedMap; +import java.util.TreeMap; +import java.util.function.Function; +import java.util.stream.Collectors; +import org.apache.arrow.util.Preconditions; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.core.udf.UdfSignature; + + +/// The result of [UdfTestFramework#execute()]. +/// +/// This class is basically a Map from Udf to [UdfTestResult#ByScenario]. +/// +/// In order to serialize and deserialize this class with Jackson, we use a [UdfTestResult.Dto] pattern where +/// instead of using [Udf], [UdfSignature], etc we use their string identifiers as keys in the map. These DTOs can be +/// created using the [#asDto] method, which converts the result into a DTO representation, and can be converted back +/// to a [UdfTestResult] using the [Dto#asUdfTestResult(Function, Function)] method, which requires functions to map +/// ids to the actual UDF and scenario objects. This means that in order to convert the DTO back to a [UdfTestResult], +/// you need to have access to the UDFs and scenarios that were used in the test, as the DTOs only contain their IDs. +/// This should not be problematic, as: +/// * The DTOs are not used to transfer the results between processes, but rather to serialize them to a file in order +/// to use it later on different tests. +/// * The DTOs are usually good enough to create reports or to display the results in a UI. +public class UdfTestResult { + + private final Map _results; + + public UdfTestResult(Map results) { + _results = results; + } + + public Map getResults() { + return _results; + } + + public Dto asDto() { + Map dtoMap = _results.entrySet().stream() + .collect(Collectors.toMap( + entry -> entry.getKey().getMainName(), + entry -> entry.getValue().asDto() + )); + return new Dto(dtoMap); + } + + public static class Dto { + private final Map _map; + + @JsonCreator + public Dto(@JsonAnySetter Map map) { + _map = map; + } + + @JsonAnyGetter + public Map getMap() { + return _map; + } + + /// Converts this DTO into a [UdfTestResult] using the provided functions to map UDF and scenario IDs to their + /// actual objects. + /// @param udfById a function that maps a UDF ID to the actual UDF object + /// @param scenarioById a function that maps a scenario ID to the actual UdfTestScenario object + public UdfTestResult asUdfTestResult( + Function udfById, + Function scenarioById + ) { + Map udfMap = Maps.newHashMapWithExpectedSize(_map.size()); + + for (Map.Entry entry : _map.entrySet()) { + String udfId = entry.getKey(); + Udf udf = udfById.apply(udfId); + Preconditions.checkState(udf != null, "No UDF object found for: %s", udfId); + ByScenario byScenario = entry.getValue().asByScenario(udf, scenarioById); + udfMap.put(udf, byScenario); + } + return new UdfTestResult(udfMap); + } + } + + public static class ByScenario { + private final Map _map; + + public ByScenario(Map map) { + _map = map; + } + + public Map getMap() { + return _map; + } + + public Dto asDto() { + TreeMap dtoMap = _map.entrySet().stream() + .collect(Collectors.toMap( + entry -> entry.getKey().getTitle(), + entry -> entry.getValue().asDto(), + (existing, replacement) -> existing, // This should not happen, but just in case + TreeMap::new + )); + return new Dto(dtoMap); + } + + public static class Dto { + private final SortedMap _map; + + @JsonCreator + public Dto(@JsonAnySetter SortedMap map) { + _map = map; + } + + @JsonAnyGetter + public SortedMap getMap() { + return _map; + } + + /// Converts this DTO into a [ByScenario] using the provided UDF and scenario functions. + /// @param udf the UDF object to use for the conversion + /// @param scenarioById a function that maps a scenario ID to the actual UdfTestScenario object + public ByScenario asByScenario(Udf udf, Function scenarioById) { + Map scenarioMap = Maps.newHashMapWithExpectedSize(_map.size()); + + for (Map.Entry entry : _map.entrySet()) { + String scenarioId = entry.getKey(); + UdfTestScenario scenario = scenarioById.apply(scenarioId); + Preconditions.checkState(scenario != null, "No scenario object found for: %s", scenarioId); + + scenarioMap.put(scenario, entry.getValue().asBySignature(udf)); + } + return new ByScenario(scenarioMap); + } + } + } + + public static class BySignature { + private final Map _map; + + public BySignature(Map map) { + _map = map; + } + + public Map getMap() { + return _map; + } + + public Dto asDto() { + SortedMap dtoMap = _map.entrySet().stream() + .collect(Collectors.toMap( + entry -> entry.getKey().toString(), + entry -> entry.getValue().asDto(), + (existing, replacement) -> existing, // This should not happen, but just in case + TreeMap::new + )); + return new Dto(dtoMap); + } + + public static class Dto { + private final SortedMap _map; + + @JsonCreator + public Dto(@JsonAnySetter SortedMap map) { + _map = map; + } + + @JsonAnyGetter + public SortedMap getMap() { + return _map; + } + + /// Converts this DTO into a [BySignature] using the provided UDF object. + /// @param udf the UDF object to use for the conversion + public BySignature asBySignature(Udf udf) { + Map signatureMap = Maps.newHashMapWithExpectedSize(_map.size()); + + Map signatureByString = udf.getExamples().keySet().stream() + .collect(Collectors.toMap( + UdfSignature::toString, + Function.identity() + )); + + for (Map.Entry entry : _map.entrySet()) { + UdfSignature signature = signatureByString.get(entry.getKey()); + + Preconditions.checkState(signature != null, "No signature object found for: %s", entry.getKey()); + Map exampleById = udf.getExamples().get(signature).stream() + .collect(Collectors.toMap( + UdfExample::getId, + Function.identity() + )); + + signatureMap.put(signature, entry.getValue().asResultByExample(exampleById::get)); + } + return new BySignature(signatureMap); + } + } + } +} diff --git a/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/UdfTestScenario.java b/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/UdfTestScenario.java new file mode 100644 index 000000000000..5eb6aacede96 --- /dev/null +++ b/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/UdfTestScenario.java @@ -0,0 +1,63 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.udf.test; + +import java.util.Map; +import java.util.function.Function; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.core.udf.UdfSignature; + +/// A scenario for testing [UDFs][Udf] (User Defined Functions). +/// +/// This interface is used to define _a way to test a UDF_ with a given set of examples. +/// This is necessary because Pinot supports a wide range of query engines and even the same engine may use UDFs in +/// different ways. For example, it is not the same to test a UDF in the context of a SSE block transform, a MSE +/// intermediate projection or a row transform during ingestion. +public interface UdfTestScenario { + + /// A title for the scenario, used in reports. + String getTitle(); + + /// A description of the scenario, used in reports. + String getDescription(); + + /// Execute the test scenario for the given UDF suite and signature. + Map execute(Udf suite, UdfSignature signature); + + /// A factory interface to create scenarios. + /// This is mainly used to be able to register these scenarios using ServiceLoader, so that they can be + /// discovered and used by the UDF test framework. + interface Factory { + UdfTestScenario create(UdfTestCluster cluster); + + class FromCluster implements Factory { + private final Function _factoryFun; + + public FromCluster(Function factoryFun) { + _factoryFun = factoryFun; + } + + @Override + public UdfTestScenario create(UdfTestCluster cluster) { + return _factoryFun.apply(cluster); + } + } + } +} diff --git a/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/scenarios/AbstractUdfTestScenario.java b/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/scenarios/AbstractUdfTestScenario.java new file mode 100644 index 000000000000..d87fe5d9d95f --- /dev/null +++ b/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/scenarios/AbstractUdfTestScenario.java @@ -0,0 +1,125 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.udf.test.scenarios; + +import com.google.common.collect.Maps; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.core.udf.UdfExample.NullHandling; +import org.apache.pinot.core.udf.UdfSignature; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.udf.test.PinotFunctionEnvGenerator; +import org.apache.pinot.udf.test.UdfExampleResult; +import org.apache.pinot.udf.test.UdfTestCluster; +import org.apache.pinot.udf.test.UdfTestScenario; + + +/// An abstract class for UDF test scenarios, providing common functionality like running SQL queries on a given +/// PinotFunctionTestCluster and extracting results from the query output. +public abstract class AbstractUdfTestScenario implements UdfTestScenario { + + protected final UdfTestCluster _cluster; + private final NullHandling _nullHandlingMode; + + public AbstractUdfTestScenario(UdfTestCluster cluster, NullHandling nullHandlingMode) { + _cluster = cluster; + _nullHandlingMode = nullHandlingMode; + } + + protected NullHandling getNullHandlingMode() { + return _nullHandlingMode; + } + + protected String replaceCommonVariables( + Udf udf, + UdfSignature signature, + NullHandling nullHandling, + /* language=sql*/ String templateSql) { + return templateSql + .replace("@table", PinotFunctionEnvGenerator.getTableName(udf)) + .replace("@udfCol", PinotFunctionEnvGenerator.getUdfColumnName()) + .replace("@udfName", udf.getMainCanonicalName()) + .replace("@testCol", PinotFunctionEnvGenerator.getTestColumnName()) + .replace("@resultCol", PinotFunctionEnvGenerator.getResultColumnName(signature, nullHandling)) + // Important, we need to replace the @testCol and @resultCol before replacing @test and @result respectively + .replace("@test", "test") + .replace("@result", "result") + .replace("@signatureCol", PinotFunctionEnvGenerator.getSignatureColumnName()) + .replace("@signature", signature.toString()); + } + + protected String replaceCall( + Udf udf, + UdfSignature signature, + UdfExample example, + /* language=sql*/ String templateSql) { + List argsForCall = PinotFunctionEnvGenerator.getArgsForCall(signature, example); + String call = udf.asSqlCall(udf.getMainCanonicalName(), argsForCall); + return templateSql + .replace("@example", example.getId()) + .replace("@call", call); + } + + @Override + public String toString() { + return getClass().getSimpleName(); + } + + protected Map extractResultsByCase( + Udf udf, + UdfSignature signature, + UdfTestCluster.ExecutionContext context, + /* language=sql*/ String sqlTemplate) { + Set examples = udf.getExamples().get(signature); + Map results = Maps.newHashMapWithExpectedSize(examples.size()); + + String sqlTemplate2 = replaceCommonVariables(udf, signature, getNullHandlingMode(), sqlTemplate); + + for (UdfExample example : examples) { + String sql = replaceCall(udf, signature, example, sqlTemplate2); + Iterator rows = _cluster.query(context, sql); + + if (!rows.hasNext()) { + String errorMessage = "No results found for example: " + example.getId(); + results.put(example, UdfExampleResult.error(example, errorMessage)); + continue; + } + GenericRow row = rows.next(); + String testId = (String) row.getValue("test"); + if (testId != null && !testId.equals(example.getId())) { + String errorMessage = "Test ID mismatch: expected " + example.getId() + ", found " + testId; + results.put(example, UdfExampleResult.error(example, errorMessage)); + continue; + } + Object expectedResult = example.getResult(getNullHandlingMode()); + results.put(example, UdfExampleResult.success(example, row.getValue("result"), expectedResult)); + + if (rows.hasNext()) { + String errorMessage = "Multiple results found for example: " + example.getId(); + results.put(example, UdfExampleResult.error(example, errorMessage)); + } + } + + return results; + } +} diff --git a/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/scenarios/ExpressionTransformerTestScenario.java b/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/scenarios/ExpressionTransformerTestScenario.java new file mode 100644 index 000000000000..5f80c5093575 --- /dev/null +++ b/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/scenarios/ExpressionTransformerTestScenario.java @@ -0,0 +1,78 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.udf.test.scenarios; + +import com.google.auto.service.AutoService; +import java.util.HashMap; +import java.util.Map; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.core.udf.UdfSignature; +import org.apache.pinot.segment.local.function.FunctionEvaluator; +import org.apache.pinot.segment.local.function.FunctionEvaluatorFactory; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.udf.test.PinotFunctionEnvGenerator; +import org.apache.pinot.udf.test.UdfExampleResult; +import org.apache.pinot.udf.test.UdfTestCluster; +import org.apache.pinot.udf.test.UdfTestScenario; + + +/// A test scenario where the UDF is executed as an ingestion time transformer. This scenario doesn't actually use +/// a cluster but instead uses the [FunctionEvaluatorFactory] to create an evaluator for the UDF and +/// evaluate it directly on the [GenericRow] objects generated by the [PinotFunctionEnvGenerator]. +public class ExpressionTransformerTestScenario implements UdfTestScenario { + + @Override + public String getTitle() { + return "Ingestion time transformer"; + } + + @Override + public String getDescription() { + return "This scenario tests the UDF as an ingestion time transformer."; + } + + @Override + public Map execute(Udf udf, UdfSignature signature) { + + Map result = new HashMap<>(); + + for (UdfExample testCase : udf.getExamples().get(signature)) { + String sqlCall = udf.asSqlCall(udf.getMainCanonicalName(), PinotFunctionEnvGenerator.getArgsForCall(signature)); + FunctionEvaluator funEvaluator = FunctionEvaluatorFactory.getExpressionEvaluator(sqlCall); + GenericRow row = PinotFunctionEnvGenerator.asRow(udf, signature, testCase); + try { + Object callResult = funEvaluator.evaluate(row); + Object expectedResult = testCase.getResult(UdfExample.NullHandling.ENABLED); + result.put(testCase, UdfExampleResult.success(testCase, callResult, expectedResult)); + } catch (Exception e) { + result.put(testCase, UdfExampleResult.error(testCase, e.getMessage())); + } + } + return result; + } + + @AutoService(UdfTestScenario.Factory.class) + public static class Factory implements UdfTestScenario.Factory { + @Override + public UdfTestScenario create(UdfTestCluster cluster) { + return new ExpressionTransformerTestScenario(); + } + } +} diff --git a/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/scenarios/IntermediateUdfTestScenario.java b/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/scenarios/IntermediateUdfTestScenario.java new file mode 100644 index 000000000000..5dbb090f4223 --- /dev/null +++ b/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/scenarios/IntermediateUdfTestScenario.java @@ -0,0 +1,130 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.udf.test.scenarios; + +import com.google.auto.service.AutoService; +import com.google.common.collect.Maps; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.core.udf.UdfParameter; +import org.apache.pinot.core.udf.UdfSignature; +import org.apache.pinot.udf.test.PinotFunctionEnvGenerator; +import org.apache.pinot.udf.test.UdfExampleResult; +import org.apache.pinot.udf.test.UdfTestCluster; +import org.apache.pinot.udf.test.UdfTestScenario; + + +/// A test scenario where the UDF is executed on an intermediate stage of the MSE. +public class IntermediateUdfTestScenario extends AbstractUdfTestScenario { + + public IntermediateUdfTestScenario(UdfTestCluster cluster, UdfExample.NullHandling nullHandlingMode) { + super(cluster, nullHandlingMode); + } + + @Override + public String getTitle() { + return "MSE intermediate stage (" + (getNullHandlingMode() == UdfExample.NullHandling.ENABLED ? "with" : "without") + + " null handling)"; + } + + @Override + public String getDescription() { + return "This scenario tests the UDF as an intermediate stage of the MSE. "; + } + + @Override + public Map execute( + Udf udf, + UdfSignature signature) { + List params = signature.getParameters(); + + Set examples = udf.getExamples().get(signature); + Map results = Maps.newHashMapWithExpectedSize(examples.size()); + // TODO: Look for a way to force the UDF to be executed on an intermediate stage of the MSE when it has 0 params. + if (params.isEmpty()) { + for (UdfExample example : examples) { + String errorMessage = "Need at least one parameter to execute the UDF in an intermediate stage"; + results.put(example, UdfExampleResult.error(example, errorMessage)); + } + return results; + } + + String firstCol = PinotFunctionEnvGenerator.getParameterColumnName(signature, 0); + StringBuilder otherCols = new StringBuilder(); + for (int i = 1; i < params.size(); i++) { + String colName = PinotFunctionEnvGenerator.getParameterColumnName(signature, i); + otherCols.append(",\n t1.") + .append(colName) + .append(" AS ") + .append(colName); + } + + // language=sql + String sqlTemplate = ("" + + "WITH fakeTable AS (\n" // this table is used to make sure the call is made on an intermediate stage + + " SELECT \n" + + " t1.@udfCol, \n" + + " t1.@signatureCol, \n" + + " t1.@testCol, \n" + // Calcite is not smart enough to know that this coalesce could be simplified to just t1.@firstCol, + // and given it uses cols from two different columns it has to be executed after the join. + + " coalesce(t1.@firstCol, t2.@firstCol) as @firstCol" // + + "@otherCols\n" + + " FROM @table_OFFLINE AS t1 \n" + + " JOIN @table_OFFLINE AS t2 \n" + + " ON t1.@testCol = t2.@testCol AND t1.@signatureCol = t2.@signatureCol \n" + + ")\n" + + "SELECT \n" + + " @testCol as test, \n" + // This call will use the result of the coalesce as argument. This means it must be executed after the join, + // which means that it has to be executed in an intermediate stage. + + " @call AS result\n" + + "FROM fakeTable \n" + + "WHERE @signatureCol = '@signature' \n" + + " AND @udfCol = '@udfName' \n" + + " AND @testCol = '@example' \n") + .replace("@firstCol", firstCol) + .replace("@otherCols", otherCols.toString()); + + UdfTestCluster.ExecutionContext context = new UdfTestCluster.ExecutionContext( + getNullHandlingMode(), + true); + + return extractResultsByCase(udf, signature, context, sqlTemplate); + } + + /// A factory that creates an instance of this scenario with null handling enabled + @AutoService(UdfTestScenario.Factory.class) + public static class WithNullHandlingFactory extends UdfTestScenario.Factory.FromCluster { + public WithNullHandlingFactory() { + super(cluster -> new IntermediateUdfTestScenario(cluster, UdfExample.NullHandling.ENABLED)); + } + } + + /// A factory that creates an instance of this scenario with null handling disabled + @AutoService(UdfTestScenario.Factory.class) + public static class WithoutNullHandlingFactory extends UdfTestScenario.Factory.FromCluster { + public WithoutNullHandlingFactory() { + super(cluster -> new IntermediateUdfTestScenario(cluster, UdfExample.NullHandling.DISABLED)); + } + } +} diff --git a/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/scenarios/PredicateUdfTestScenario.java b/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/scenarios/PredicateUdfTestScenario.java new file mode 100644 index 000000000000..c931d24534ee --- /dev/null +++ b/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/scenarios/PredicateUdfTestScenario.java @@ -0,0 +1,117 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.udf.test.scenarios; + +import com.google.auto.service.AutoService; +import com.google.common.collect.Sets; +import java.util.HashMap; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.core.udf.UdfSignature; +import org.apache.pinot.udf.test.UdfExampleResult; +import org.apache.pinot.udf.test.UdfTestCluster; +import org.apache.pinot.udf.test.UdfTestScenario; + + +/// A test scenario where the UDF is executed as a predicate, which is always evaluated as a ScalarFunction. +public class PredicateUdfTestScenario extends AbstractUdfTestScenario { + public PredicateUdfTestScenario(UdfTestCluster cluster, UdfExample.NullHandling nullHandlingMode) { + super(cluster, nullHandlingMode); + } + + @Override + public String getTitle() { + return "SSE predicate (" + (getNullHandlingMode() == UdfExample.NullHandling.ENABLED ? "with" : "without") + + " null handling)"; + } + + @Override + public String getDescription() { + return "This scenario tests the UDF as a predicate in the SSE."; + } + + @Override + public Map execute(Udf udf, UdfSignature signature) { + + // language=sql + String sqlTemplate = "" + + "SELECT \n" + + " @testCol as test," + + " true as result\n" + + "FROM @table \n" + + "WHERE @signatureCol = '@signature' \n" + + " AND @udfCol = '@udfName' \n" + + " AND ( \n" + + " @call = @resultCol OR \n" + + " @call IS NULL AND @resultCol IS NULL \n" + + " )\n" + + " AND @testCol = '@example' \n"; + UdfTestCluster.ExecutionContext context = new UdfTestCluster.ExecutionContext( + getNullHandlingMode(), + false); + Map queryResult = extractResultsByCase(udf, signature, context, sqlTemplate); + + Set successfulCases = queryResult.values() + .stream() + .map(UdfExampleResult::getTest) + .collect(Collectors.toSet()); + + Set examples = udf.getExamples().get(signature); + Map result = new HashMap<>(); + Sets.SetView expectedPositives = Sets.intersection(examples, successfulCases); + try { + for (UdfExample expectedPositive : expectedPositives) { + result.put(expectedPositive, UdfExampleResult.success(expectedPositive, true, true)); + } + } catch (RuntimeException e) { + throw e; + } + + if (expectedPositives.size() != udf.getExamples().size()) { + Sets.SetView unexpectedNegatives = Sets.difference(examples, successfulCases); + for (UdfExample unexpected : unexpectedNegatives) { + result.put(unexpected, UdfExampleResult.success(unexpected, false, true)); + } + Sets.SetView unexpectedPositives = Sets.difference(successfulCases, examples); + for (UdfExample unexpected : unexpectedPositives) { + result.put(unexpected, UdfExampleResult.success(unexpected, true, false)); + } + } + return result; + } + + /// A factory that creates an instance of this scenario with null handling enabled + @AutoService(UdfTestScenario.Factory.class) + public static class WithNullHandlingFactory extends UdfTestScenario.Factory.FromCluster { + public WithNullHandlingFactory() { + super(cluster -> new PredicateUdfTestScenario(cluster, UdfExample.NullHandling.ENABLED)); + } + } + + /// A factory that creates an instance of this scenario with null handling disabled + @AutoService(UdfTestScenario.Factory.class) + public static class WithoutNullHandlingFactory extends UdfTestScenario.Factory.FromCluster { + public WithoutNullHandlingFactory() { + super(cluster -> new PredicateUdfTestScenario(cluster, UdfExample.NullHandling.DISABLED)); + } + } +} diff --git a/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/scenarios/TransformationUdfTestScenario.java b/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/scenarios/TransformationUdfTestScenario.java new file mode 100644 index 000000000000..86bced3ed967 --- /dev/null +++ b/pinot-udf-test/src/main/java/org/apache/pinot/udf/test/scenarios/TransformationUdfTestScenario.java @@ -0,0 +1,85 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.udf.test.scenarios; + +import com.google.auto.service.AutoService; +import java.util.Map; +import org.apache.pinot.core.udf.Udf; +import org.apache.pinot.core.udf.UdfExample; +import org.apache.pinot.core.udf.UdfSignature; +import org.apache.pinot.udf.test.UdfExampleResult; +import org.apache.pinot.udf.test.UdfTestCluster; +import org.apache.pinot.udf.test.UdfTestScenario; + + +/// A test scenario where the UDF is executed as a TransformFunction. +public class TransformationUdfTestScenario extends AbstractUdfTestScenario { + public TransformationUdfTestScenario(UdfTestCluster cluster, UdfExample.NullHandling nullHandlingMode) { + super(cluster, nullHandlingMode); + } + + @Override + public String getTitle() { + return "SSE projection (" + (getNullHandlingMode() == UdfExample.NullHandling.ENABLED ? "with" : "without") + + " null handling)"; + } + + @Override + public String getDescription() { + return "This scenario tests the UDF as a projection in the SSE."; + } + + @Override + public Map execute( + Udf suite, + UdfSignature signature) { + + // language=sql + String sqlTemplate = "" + + "SELECT \n" + + " @testCol as test, \n" + + " @call AS result\n" + + "FROM @table \n" + + "WHERE @signatureCol = '@signature' \n" + + " AND @udfCol = '@udfName' \n" + + " AND @testCol = '@example' \n"; + + UdfTestCluster.ExecutionContext context = new UdfTestCluster.ExecutionContext( + getNullHandlingMode(), + false); + + return extractResultsByCase(suite, signature, context, sqlTemplate); + } + + /// A factory that creates an instance of this scenario with null handling enabled + @AutoService(UdfTestScenario.Factory.class) + public static class WithNullHandlingFactory extends UdfTestScenario.Factory.FromCluster { + public WithNullHandlingFactory() { + super(cluster -> new TransformationUdfTestScenario(cluster, UdfExample.NullHandling.ENABLED)); + } + } + + /// A factory that creates an instance of this scenario with null handling disabled + @AutoService(UdfTestScenario.Factory.class) + public static class WithoutNullHandlingFactory extends UdfTestScenario.Factory.FromCluster { + public WithoutNullHandlingFactory() { + super(cluster -> new TransformationUdfTestScenario(cluster, UdfExample.NullHandling.DISABLED)); + } + } +} diff --git a/pom.xml b/pom.xml index ccccac7f5c3f..d2024d480647 100644 --- a/pom.xml +++ b/pom.xml @@ -60,6 +60,7 @@ pinot-query-planner pinot-query-runtime pinot-timeseries + pinot-udf-test @@ -628,6 +629,11 @@ pinot-timeseries-m3ql ${project.version} + + org.apache.pinot + pinot-udf-test + ${project.version} + From 0415a2bc750cc71d919b7f54bc325f4b7a9e3b4b Mon Sep 17 00:00:00 2001 From: RAGHVENDRA KUMAR YADAV Date: Thu, 24 Jul 2025 04:07:58 -0700 Subject: [PATCH 015/167] Make ForwardIndexReader pluggable by removing static methods from ForwardIndexType and ForwardIndexReaderFactory (#16363) --- .../forward/ForwardIndexReaderFactory.java | 6 +-- .../index/forward/ForwardIndexType.java | 35 -------------- .../index/loader/ForwardIndexHandler.java | 33 ++++++++++--- .../bloomfilter/BloomFilterHandler.java | 6 ++- .../ColumnMinMaxValueGenerator.java | 3 +- .../BaseDefaultColumnHandler.java | 12 +++-- .../loader/invertedindex/H3IndexHandler.java | 6 ++- .../invertedindex/InvertedIndexHandler.java | 6 ++- .../invertedindex/JsonIndexHandler.java | 10 ++-- .../MultiColumnTextIndexHandler.java | 9 ++-- .../invertedindex/RangeIndexHandler.java | 10 ++-- .../invertedindex/TextIndexHandler.java | 7 +-- .../invertedindex/VectorIndexHandler.java | 6 ++- .../v2/store/StarTreeLoaderUtils.java | 4 +- ...ultiValueFixedByteRawIndexCreatorTest.java | 3 +- .../MultiValueVarByteRawIndexCreatorTest.java | 6 ++- .../index/creator/RawIndexCreatorTest.java | 11 +++-- .../index/loader/ForwardIndexHandlerTest.java | 48 +++++++++++++++---- .../index/loader/SegmentPreProcessorTest.java | 12 +++-- 19 files changed, 142 insertions(+), 91 deletions(-) diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/forward/ForwardIndexReaderFactory.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/forward/ForwardIndexReaderFactory.java index 59a69047c0df..f7bf16ca93e6 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/forward/ForwardIndexReaderFactory.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/forward/ForwardIndexReaderFactory.java @@ -71,7 +71,7 @@ protected ForwardIndexReader createIndexReader(PinotDataBuffer dataBuffer, Colum return createIndexReader(dataBuffer, metadata); } - public static ForwardIndexReader createIndexReader(PinotDataBuffer dataBuffer, ColumnMetadata metadata) { + public ForwardIndexReader createIndexReader(PinotDataBuffer dataBuffer, ColumnMetadata metadata) { if (metadata.hasDictionary()) { if (metadata.isSingleValue()) { if (metadata.isSorted()) { @@ -108,7 +108,7 @@ public static ForwardIndexReader createIndexReader(PinotDataBuffer dataBuffer, C } } - public static ForwardIndexReader createRawIndexReader(PinotDataBuffer dataBuffer, DataType storedType, + public ForwardIndexReader createRawIndexReader(PinotDataBuffer dataBuffer, DataType storedType, boolean isSingleValue) { int version = dataBuffer.getInt(0); if (isSingleValue && storedType.isFixedWidth()) { @@ -128,7 +128,7 @@ public static ForwardIndexReader createRawIndexReader(PinotDataBuffer dataBuffer } } - private static ForwardIndexReader createNonV4RawIndexReader(PinotDataBuffer dataBuffer, DataType storedType, + private ForwardIndexReader createNonV4RawIndexReader(PinotDataBuffer dataBuffer, DataType storedType, boolean isSingleValue) { // Only reach here if SV + raw + var byte + non v4 or MV + non v4 if (isSingleValue) { diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/forward/ForwardIndexType.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/forward/ForwardIndexType.java index 0c95bf426bd3..4e5ac20a3ef4 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/forward/ForwardIndexType.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/forward/ForwardIndexType.java @@ -21,7 +21,6 @@ import com.google.common.base.Preconditions; import com.google.common.collect.Maps; -import java.io.IOException; import java.util.Collection; import java.util.Collections; import java.util.List; @@ -43,7 +42,6 @@ import org.apache.pinot.segment.spi.index.FieldIndexConfigs; import org.apache.pinot.segment.spi.index.ForwardIndexConfig; import org.apache.pinot.segment.spi.index.IndexHandler; -import org.apache.pinot.segment.spi.index.IndexReaderConstraintException; import org.apache.pinot.segment.spi.index.IndexReaderFactory; import org.apache.pinot.segment.spi.index.IndexUtil; import org.apache.pinot.segment.spi.index.RangeIndexConfig; @@ -52,7 +50,6 @@ import org.apache.pinot.segment.spi.index.mutable.MutableIndex; import org.apache.pinot.segment.spi.index.mutable.provider.MutableIndexContext; import org.apache.pinot.segment.spi.index.reader.ForwardIndexReader; -import org.apache.pinot.segment.spi.memory.PinotDataBuffer; import org.apache.pinot.segment.spi.store.SegmentDirectory; import org.apache.pinot.spi.config.table.FieldConfig; import org.apache.pinot.spi.config.table.FieldConfig.CompressionCodec; @@ -262,38 +259,6 @@ public List getFileExtensions(@Nullable ColumnMetadata columnMetadata) { return Collections.singletonList(getFileExtension(columnMetadata)); } - /** - * Returns the forward index reader for the given column. - * - * This method will return the default reader, skipping any index overload. - */ - public static ForwardIndexReader read(SegmentDirectory.Reader segmentReader, ColumnMetadata columnMetadata) - throws IOException { - PinotDataBuffer dataBuffer = segmentReader.getIndexFor(columnMetadata.getColumnName(), StandardIndexes.forward()); - return read(dataBuffer, columnMetadata); - } - - /** - * Returns the forward index reader for the given column. - * - * This method will return the default reader, skipping any index overload. - */ - public static ForwardIndexReader read(PinotDataBuffer dataBuffer, ColumnMetadata metadata) { - return ForwardIndexReaderFactory.createIndexReader(dataBuffer, metadata); - } - - /** - * Returns the forward index reader for the given column. - * - * This method will delegate on {@link StandardIndexes}, so the correct reader will be returned even when using - * index overload. - */ - public static ForwardIndexReader read(SegmentDirectory.Reader segmentReader, FieldIndexConfigs fieldIndexConfigs, - ColumnMetadata metadata) - throws IndexReaderConstraintException, IOException { - return StandardIndexes.forward().getReaderFactory().createIndexReader(segmentReader, fieldIndexConfigs, metadata); - } - @Nullable @Override public MutableIndex createMutableIndex(MutableIndexContext context, ForwardIndexConfig config) { diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/ForwardIndexHandler.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/ForwardIndexHandler.java index c86275b43667..f63e71f54e0d 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/ForwardIndexHandler.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/ForwardIndexHandler.java @@ -43,7 +43,6 @@ import org.apache.pinot.segment.local.segment.creator.impl.stats.MapColumnPreIndexStatsCollector; import org.apache.pinot.segment.local.segment.creator.impl.stats.StringColumnPreIndexStatsCollector; import org.apache.pinot.segment.local.segment.index.dictionary.DictionaryIndexType; -import org.apache.pinot.segment.local.segment.index.forward.ForwardIndexType; import org.apache.pinot.segment.local.segment.readers.PinotSegmentColumnReader; import org.apache.pinot.segment.spi.ColumnMetadata; import org.apache.pinot.segment.spi.V1Constants; @@ -56,6 +55,7 @@ import org.apache.pinot.segment.spi.index.FieldIndexConfigs; import org.apache.pinot.segment.spi.index.FieldIndexConfigsUtil; import org.apache.pinot.segment.spi.index.ForwardIndexConfig; +import org.apache.pinot.segment.spi.index.IndexReaderFactory; import org.apache.pinot.segment.spi.index.IndexType; import org.apache.pinot.segment.spi.index.StandardIndexes; import org.apache.pinot.segment.spi.index.creator.ForwardIndexCreator; @@ -376,7 +376,11 @@ private boolean shouldChangeRawCompressionType(String column, SegmentDirectory.R // The compression type for an existing segment can only be determined by reading the forward index header. ColumnMetadata existingColMetadata = _segmentDirectory.getSegmentMetadata().getColumnMetadataFor(column); ChunkCompressionType existingCompressionType; - try (ForwardIndexReader fwdIndexReader = ForwardIndexType.read(segmentReader, existingColMetadata)) { + + // Get the forward index reader factory and create a reader + IndexReaderFactory readerFactory = StandardIndexes.forward().getReaderFactory(); + try (ForwardIndexReader fwdIndexReader = readerFactory.createIndexReader(segmentReader, + _fieldIndexConfigs.get(column), existingColMetadata)) { existingCompressionType = fwdIndexReader.getCompressionType(); Preconditions.checkState(existingCompressionType != null, "Existing compressionType cannot be null for raw forward index column=" + column); @@ -397,7 +401,10 @@ private boolean shouldChangeDictIdCompressionType(String column, SegmentDirector // The compression type for an existing segment can only be determined by reading the forward index header. ColumnMetadata existingColMetadata = _segmentDirectory.getSegmentMetadata().getColumnMetadataFor(column); DictIdCompressionType existingCompressionType; - try (ForwardIndexReader fwdIndexReader = ForwardIndexType.read(segmentReader, existingColMetadata)) { + // Get the forward index reader factory and create a reader + IndexReaderFactory readerFactory = StandardIndexes.forward().getReaderFactory(); + try (ForwardIndexReader fwdIndexReader = readerFactory.createIndexReader(segmentReader, + _fieldIndexConfigs.get(column), existingColMetadata)) { existingCompressionType = fwdIndexReader.getDictIdCompressionType(); } @@ -460,7 +467,10 @@ private void rewriteForwardIndexForCompressionChange(String column, SegmentDirec private void rewriteForwardIndexForCompressionChange(String column, ColumnMetadata columnMetadata, File indexDir, SegmentDirectory.Writer segmentWriter) throws Exception { - try (ForwardIndexReader reader = ForwardIndexType.read(segmentWriter, columnMetadata)) { + // Get the forward index reader factory and create a reader + IndexReaderFactory readerFactory = StandardIndexes.forward().getReaderFactory(); + try (ForwardIndexReader reader = readerFactory.createIndexReader(segmentWriter, _fieldIndexConfigs.get(column), + columnMetadata)) { IndexCreationContext.Builder builder = IndexCreationContext.builder().withIndexDir(indexDir).withColumnMetadata(columnMetadata) .withTableNameWithType(_tableConfig.getTableName()) @@ -859,7 +869,10 @@ private SegmentDictionaryCreator buildDictionary(String column, ColumnMetadata e throws Exception { int numDocs = existingColMetadata.getTotalDocs(); - try (ForwardIndexReader reader = ForwardIndexType.read(segmentWriter, existingColMetadata)) { + // Get the forward index reader factory and create a reader + IndexReaderFactory readerFactory = StandardIndexes.forward().getReaderFactory(); + try (ForwardIndexReader reader = readerFactory.createIndexReader(segmentWriter, _fieldIndexConfigs.get(column), + existingColMetadata)) { // Note: Special Null handling is not necessary here. This is because, the existing default null value in the // raw forwardIndex will be retained as such while created the dictionary and dict-based forward index. Also, // null value vectors maintain a bitmap of docIds. No handling is necessary there. @@ -888,7 +901,10 @@ private SegmentDictionaryCreator buildDictionary(String column, ColumnMetadata e private void writeDictEnabledForwardIndex(String column, ColumnMetadata existingColMetadata, SegmentDirectory.Writer segmentWriter, File indexDir, SegmentDictionaryCreator dictionaryCreator) throws Exception { - try (ForwardIndexReader reader = ForwardIndexType.read(segmentWriter, existingColMetadata)) { + // Get the forward index reader factory and create a reader + IndexReaderFactory readerFactory = StandardIndexes.forward().getReaderFactory(); + try (ForwardIndexReader reader = readerFactory.createIndexReader(segmentWriter, _fieldIndexConfigs.get(column), + existingColMetadata)) { IndexCreationContext.Builder builder = IndexCreationContext.builder().withIndexDir(indexDir).withColumnMetadata(existingColMetadata) .withTableNameWithType(_tableConfig.getTableName()) @@ -964,7 +980,10 @@ private void rewriteDictToRawForwardIndex(ColumnMetadata columnMetadata, Segment File indexDir) throws Exception { String column = columnMetadata.getColumnName(); - try (ForwardIndexReader reader = ForwardIndexType.read(segmentWriter, columnMetadata)) { + // Get the forward index reader factory and create a reader + IndexReaderFactory readerFactory = StandardIndexes.forward().getReaderFactory(); + try (ForwardIndexReader reader = readerFactory.createIndexReader(segmentWriter, _fieldIndexConfigs.get(column), + columnMetadata)) { Dictionary dictionary = DictionaryIndexType.read(segmentWriter, columnMetadata); IndexCreationContext.Builder builder = IndexCreationContext.builder().withIndexDir(indexDir).withColumnMetadata(columnMetadata) diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/bloomfilter/BloomFilterHandler.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/bloomfilter/BloomFilterHandler.java index 95500a60cd51..7c05a3fa4c5d 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/bloomfilter/BloomFilterHandler.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/bloomfilter/BloomFilterHandler.java @@ -25,7 +25,6 @@ import java.util.Set; import org.apache.commons.io.FileUtils; import org.apache.pinot.segment.local.segment.index.dictionary.DictionaryIndexType; -import org.apache.pinot.segment.local.segment.index.forward.ForwardIndexType; import org.apache.pinot.segment.local.segment.index.loader.BaseIndexHandler; import org.apache.pinot.segment.local.segment.index.loader.LoaderUtils; import org.apache.pinot.segment.spi.ColumnMetadata; @@ -35,6 +34,7 @@ import org.apache.pinot.segment.spi.index.DictionaryIndexConfig; import org.apache.pinot.segment.spi.index.FieldIndexConfigs; import org.apache.pinot.segment.spi.index.FieldIndexConfigsUtil; +import org.apache.pinot.segment.spi.index.IndexReaderFactory; import org.apache.pinot.segment.spi.index.StandardIndexes; import org.apache.pinot.segment.spi.index.creator.BloomFilterCreator; import org.apache.pinot.segment.spi.index.reader.Dictionary; @@ -143,9 +143,11 @@ private void createAndSealBloomFilterForNonDictionaryColumn(File indexDir, Colum .withContinueOnError(_tableConfig.getIngestionConfig() != null && _tableConfig.getIngestionConfig().isContinueOnError()) .build(); + IndexReaderFactory readerFactory = StandardIndexes.forward().getReaderFactory(); try (BloomFilterCreator bloomFilterCreator = StandardIndexes.bloomFilter() .createIndexCreator(context, bloomFilterConfig); - ForwardIndexReader forwardIndexReader = ForwardIndexType.read(segmentWriter, columnMetadata); + ForwardIndexReader forwardIndexReader = readerFactory.createIndexReader(segmentWriter, + _fieldIndexConfigs.get(columnMetadata.getColumnName()), columnMetadata); ForwardIndexReaderContext readerContext = forwardIndexReader.createContext()) { if (columnMetadata.isSingleValue()) { // SV diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/columnminmaxvalue/ColumnMinMaxValueGenerator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/columnminmaxvalue/ColumnMinMaxValueGenerator.java index 312a6d2fac4b..d965a89835f2 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/columnminmaxvalue/ColumnMinMaxValueGenerator.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/columnminmaxvalue/ColumnMinMaxValueGenerator.java @@ -204,7 +204,8 @@ private void addColumnMinMaxValueWithoutDictionary(ColumnMetadata columnMetadata DataType storedType = dataType.getStoredType(); boolean isSingleValue = columnMetadata.isSingleValue(); PinotDataBuffer rawIndexBuffer = _segmentWriter.getIndexFor(columnName, StandardIndexes.forward()); - try (ForwardIndexReader rawIndexReader = ForwardIndexReaderFactory.createRawIndexReader(rawIndexBuffer, storedType, + try (ForwardIndexReader rawIndexReader = ForwardIndexReaderFactory.getInstance() + .createRawIndexReader(rawIndexBuffer, storedType, isSingleValue); ForwardIndexReaderContext readerContext = rawIndexReader.createContext()) { int numDocs = columnMetadata.getTotalDocs(); Object minValue; diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/defaultcolumn/BaseDefaultColumnHandler.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/defaultcolumn/BaseDefaultColumnHandler.java index 83fe886a249c..45bb854516ba 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/defaultcolumn/BaseDefaultColumnHandler.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/defaultcolumn/BaseDefaultColumnHandler.java @@ -50,7 +50,6 @@ import org.apache.pinot.segment.local.segment.index.dictionary.DictionaryIndexType; import org.apache.pinot.segment.local.segment.index.forward.ForwardIndexCreatorFactory; import org.apache.pinot.segment.local.segment.index.forward.ForwardIndexPlugin; -import org.apache.pinot.segment.local.segment.index.forward.ForwardIndexType; import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig; import org.apache.pinot.segment.local.segment.index.loader.LoaderUtils; import org.apache.pinot.segment.local.segment.readers.PinotSegmentColumnReader; @@ -63,6 +62,8 @@ import org.apache.pinot.segment.spi.index.DictionaryIndexConfig; import org.apache.pinot.segment.spi.index.FieldIndexConfigs; import org.apache.pinot.segment.spi.index.ForwardIndexConfig; +import org.apache.pinot.segment.spi.index.IndexReaderConstraintException; +import org.apache.pinot.segment.spi.index.IndexReaderFactory; import org.apache.pinot.segment.spi.index.IndexService; import org.apache.pinot.segment.spi.index.IndexType; import org.apache.pinot.segment.spi.index.StandardIndexes; @@ -1223,8 +1224,13 @@ private class ValueReader implements Closeable { final PinotSegmentColumnReader _columnReader; ValueReader(ColumnMetadata columnMetadata) - throws IOException { - _forwardIndexReader = ForwardIndexType.read(_segmentWriter, columnMetadata); + throws IOException, IndexReaderConstraintException { + + IndexReaderFactory readerFactory = StandardIndexes.forward().getReaderFactory(); + FieldIndexConfigs fieldIndexConfigs = new FieldIndexConfigs.Builder() + .add(StandardIndexes.forward(), ForwardIndexConfig.getDefault()) + .build(); + _forwardIndexReader = readerFactory.createIndexReader(_segmentWriter, fieldIndexConfigs, columnMetadata); if (columnMetadata.hasDictionary()) { _dictionary = DictionaryIndexType.read(_segmentWriter, columnMetadata); } else { diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/H3IndexHandler.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/H3IndexHandler.java index 803b0d925143..67b9a561289f 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/H3IndexHandler.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/H3IndexHandler.java @@ -24,7 +24,6 @@ import java.util.Map; import java.util.Set; import org.apache.commons.io.FileUtils; -import org.apache.pinot.segment.local.segment.index.forward.ForwardIndexType; import org.apache.pinot.segment.local.segment.index.loader.BaseIndexHandler; import org.apache.pinot.segment.local.segment.index.loader.LoaderUtils; import org.apache.pinot.segment.local.utils.GeometrySerializer; @@ -34,6 +33,7 @@ import org.apache.pinot.segment.spi.creator.SegmentVersion; import org.apache.pinot.segment.spi.index.FieldIndexConfigs; import org.apache.pinot.segment.spi.index.FieldIndexConfigsUtil; +import org.apache.pinot.segment.spi.index.IndexReaderFactory; import org.apache.pinot.segment.spi.index.StandardIndexes; import org.apache.pinot.segment.spi.index.creator.GeoSpatialIndexCreator; import org.apache.pinot.segment.spi.index.creator.H3IndexConfig; @@ -195,7 +195,9 @@ private void handleNonDictionaryBasedColumn(SegmentDirectory.Writer segmentWrite && _tableConfig.getIngestionConfig().isContinueOnError()) .build(); H3IndexConfig config = _fieldIndexConfigs.get(columnName).getConfig(StandardIndexes.h3()); - try (ForwardIndexReader forwardIndexReader = ForwardIndexType.read(segmentWriter, columnMetadata); + IndexReaderFactory readerFactory = StandardIndexes.forward().getReaderFactory(); + try (ForwardIndexReader forwardIndexReader = readerFactory.createIndexReader(segmentWriter, + _fieldIndexConfigs.get(columnMetadata.getColumnName()), columnMetadata); ForwardIndexReaderContext readerContext = forwardIndexReader.createContext(); GeoSpatialIndexCreator h3IndexCreator = StandardIndexes.h3().createIndexCreator(context, config)) { int numDocs = columnMetadata.getTotalDocs(); diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/InvertedIndexHandler.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/InvertedIndexHandler.java index 3999a136511f..ff4029b52d15 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/InvertedIndexHandler.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/InvertedIndexHandler.java @@ -23,7 +23,6 @@ import java.util.Map; import java.util.Set; import org.apache.commons.io.FileUtils; -import org.apache.pinot.segment.local.segment.index.forward.ForwardIndexType; import org.apache.pinot.segment.local.segment.index.loader.BaseIndexHandler; import org.apache.pinot.segment.local.segment.index.loader.LoaderUtils; import org.apache.pinot.segment.spi.ColumnMetadata; @@ -32,6 +31,7 @@ import org.apache.pinot.segment.spi.creator.SegmentVersion; import org.apache.pinot.segment.spi.index.FieldIndexConfigs; import org.apache.pinot.segment.spi.index.FieldIndexConfigsUtil; +import org.apache.pinot.segment.spi.index.IndexReaderFactory; import org.apache.pinot.segment.spi.index.StandardIndexes; import org.apache.pinot.segment.spi.index.creator.DictionaryBasedInvertedIndexCreator; import org.apache.pinot.segment.spi.index.reader.ForwardIndexReader; @@ -146,7 +146,9 @@ private void createInvertedIndexForColumn(SegmentDirectory.Writer segmentWriter, try (DictionaryBasedInvertedIndexCreator creator = StandardIndexes.inverted() .createIndexCreator(context, IndexConfig.ENABLED)) { - try (ForwardIndexReader forwardIndexReader = ForwardIndexType.read(segmentWriter, columnMetadata); + IndexReaderFactory readerFactory = StandardIndexes.forward().getReaderFactory(); + try (ForwardIndexReader forwardIndexReader = readerFactory.createIndexReader(segmentWriter, + _fieldIndexConfigs.get(columnMetadata.getColumnName()), columnMetadata); ForwardIndexReaderContext readerContext = forwardIndexReader.createContext()) { if (columnMetadata.isSingleValue()) { // Single-value column. diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/JsonIndexHandler.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/JsonIndexHandler.java index 2722d3a4a776..5636e447f1e1 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/JsonIndexHandler.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/JsonIndexHandler.java @@ -25,7 +25,6 @@ import java.util.Set; import org.apache.commons.io.FileUtils; import org.apache.pinot.segment.local.segment.index.dictionary.DictionaryIndexType; -import org.apache.pinot.segment.local.segment.index.forward.ForwardIndexType; import org.apache.pinot.segment.local.segment.index.loader.BaseIndexHandler; import org.apache.pinot.segment.local.segment.index.loader.LoaderUtils; import org.apache.pinot.segment.spi.ColumnMetadata; @@ -34,6 +33,7 @@ import org.apache.pinot.segment.spi.creator.SegmentVersion; import org.apache.pinot.segment.spi.index.FieldIndexConfigs; import org.apache.pinot.segment.spi.index.FieldIndexConfigsUtil; +import org.apache.pinot.segment.spi.index.IndexReaderFactory; import org.apache.pinot.segment.spi.index.StandardIndexes; import org.apache.pinot.segment.spi.index.creator.JsonIndexCreator; import org.apache.pinot.segment.spi.index.reader.Dictionary; @@ -165,7 +165,9 @@ private void handleDictionaryBasedColumn(SegmentDirectory.Writer segmentWriter, && _tableConfig.getIngestionConfig().isContinueOnError()) .build(); JsonIndexConfig config = _jsonIndexConfigs.get(columnName); - try (ForwardIndexReader forwardIndexReader = ForwardIndexType.read(segmentWriter, columnMetadata); + IndexReaderFactory readerFactory = StandardIndexes.forward().getReaderFactory(); + try (ForwardIndexReader forwardIndexReader = readerFactory.createIndexReader(segmentWriter, + _fieldIndexConfigs.get(columnMetadata.getColumnName()), columnMetadata); ForwardIndexReaderContext readerContext = forwardIndexReader.createContext(); Dictionary dictionary = DictionaryIndexType.read(segmentWriter, columnMetadata); JsonIndexCreator jsonIndexCreator = StandardIndexes.json().createIndexCreator(context, config)) { @@ -190,7 +192,9 @@ private void handleNonDictionaryBasedColumn(SegmentDirectory.Writer segmentWrite && _tableConfig.getIngestionConfig().isContinueOnError()) .build(); JsonIndexConfig config = _jsonIndexConfigs.get(columnName); - try (ForwardIndexReader forwardIndexReader = ForwardIndexType.read(segmentWriter, columnMetadata); + IndexReaderFactory readerFactory = StandardIndexes.forward().getReaderFactory(); + try (ForwardIndexReader forwardIndexReader = readerFactory.createIndexReader(segmentWriter, + _fieldIndexConfigs.get(columnMetadata.getColumnName()), columnMetadata); ForwardIndexReaderContext readerContext = forwardIndexReader.createContext(); JsonIndexCreator jsonIndexCreator = StandardIndexes.json().createIndexCreator(context, config)) { int numDocs = columnMetadata.getTotalDocs(); diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/MultiColumnTextIndexHandler.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/MultiColumnTextIndexHandler.java index 00d1f56394db..1ebc2469728d 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/MultiColumnTextIndexHandler.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/MultiColumnTextIndexHandler.java @@ -28,12 +28,13 @@ import org.apache.commons.configuration2.PropertiesConfiguration; import org.apache.pinot.segment.local.segment.creator.impl.text.MultiColumnLuceneTextIndexCreator; import org.apache.pinot.segment.local.segment.index.dictionary.DictionaryIndexType; -import org.apache.pinot.segment.local.segment.index.forward.ForwardIndexType; import org.apache.pinot.segment.local.segment.index.loader.BaseIndexHandler; import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig; import org.apache.pinot.segment.local.segment.index.loader.SegmentPreProcessor; import org.apache.pinot.segment.local.utils.TableConfigUtils; import org.apache.pinot.segment.spi.ColumnMetadata; +import org.apache.pinot.segment.spi.index.IndexReaderConstraintException; +import org.apache.pinot.segment.spi.index.IndexReaderFactory; import org.apache.pinot.segment.spi.index.StandardIndexes; import org.apache.pinot.segment.spi.index.TextIndexConfig; import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl; @@ -183,7 +184,7 @@ public void updateIndices(SegmentDirectory.Writer segmentWriter) private void createMultiColumnTextIndex(SegmentDirectory.Writer segmentWriter, MultiColumnLuceneTextIndexCreator textIndexCreator) - throws IOException { + throws IOException, IndexReaderConstraintException { SegmentMetadataImpl segmentMetadata = _segmentDirectory.getSegmentMetadata(); String segmentName = segmentMetadata.getName(); int numDocs = segmentMetadata.getTotalDocs(); @@ -201,7 +202,9 @@ private void createMultiColumnTextIndex(SegmentDirectory.Writer segmentWriter, for (int i = 0, n = indexedColumns.size(); i < n; i++) { String columnName = indexedColumns.get(i); ColumnMetadata metadata = createForwardIndexIfNeeded(segmentWriter, columnName, true); - ForwardIndexReader fwdReader = ForwardIndexType.read(segmentWriter, metadata); + IndexReaderFactory readerFactory = StandardIndexes.forward().getReaderFactory(); + ForwardIndexReader fwdReader = + readerFactory.createIndexReader(segmentWriter, _fieldIndexConfigs.get(metadata.getColumnName()), metadata); fwdReaders.add(fwdReader); fwdReaderContexts.add(fwdReader.createContext()); diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/RangeIndexHandler.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/RangeIndexHandler.java index dc1b51c920ba..1af29b870672 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/RangeIndexHandler.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/RangeIndexHandler.java @@ -24,7 +24,6 @@ import java.util.Map; import java.util.Set; import org.apache.commons.io.FileUtils; -import org.apache.pinot.segment.local.segment.index.forward.ForwardIndexType; import org.apache.pinot.segment.local.segment.index.loader.BaseIndexHandler; import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig; import org.apache.pinot.segment.local.segment.index.loader.LoaderUtils; @@ -34,6 +33,7 @@ import org.apache.pinot.segment.spi.creator.SegmentVersion; import org.apache.pinot.segment.spi.index.FieldIndexConfigs; import org.apache.pinot.segment.spi.index.FieldIndexConfigsUtil; +import org.apache.pinot.segment.spi.index.IndexReaderFactory; import org.apache.pinot.segment.spi.index.RangeIndexConfig; import org.apache.pinot.segment.spi.index.StandardIndexes; import org.apache.pinot.segment.spi.index.creator.CombinedInvertedIndexCreator; @@ -162,7 +162,9 @@ private void createRangeIndexForColumn(SegmentDirectory.Writer segmentWriter, Co private void handleDictionaryBasedColumn(SegmentDirectory.Writer segmentWriter, ColumnMetadata columnMetadata) throws Exception { int numDocs = columnMetadata.getTotalDocs(); - try (ForwardIndexReader forwardIndexReader = ForwardIndexType.read(segmentWriter, columnMetadata); + IndexReaderFactory readerFactory = StandardIndexes.forward().getReaderFactory(); + try (ForwardIndexReader forwardIndexReader = readerFactory.createIndexReader(segmentWriter, + _fieldIndexConfigs.get(columnMetadata.getColumnName()), columnMetadata); ForwardIndexReaderContext readerContext = forwardIndexReader.createContext(); CombinedInvertedIndexCreator rangeIndexCreator = newRangeIndexCreator(columnMetadata)) { if (columnMetadata.isSingleValue()) { @@ -185,7 +187,9 @@ private void handleDictionaryBasedColumn(SegmentDirectory.Writer segmentWriter, private void handleNonDictionaryBasedColumn(SegmentDirectory.Writer segmentWriter, ColumnMetadata columnMetadata) throws Exception { int numDocs = columnMetadata.getTotalDocs(); - try (ForwardIndexReader forwardIndexReader = ForwardIndexType.read(segmentWriter, columnMetadata); + IndexReaderFactory readerFactory = StandardIndexes.forward().getReaderFactory(); + try (ForwardIndexReader forwardIndexReader = readerFactory.createIndexReader(segmentWriter, + _fieldIndexConfigs.get(columnMetadata.getColumnName()), columnMetadata); ForwardIndexReaderContext readerContext = forwardIndexReader.createContext(); CombinedInvertedIndexCreator rangeIndexCreator = newRangeIndexCreator(columnMetadata)) { if (columnMetadata.isSingleValue()) { diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/TextIndexHandler.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/TextIndexHandler.java index 0033b59b3427..8d78f52e3e9a 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/TextIndexHandler.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/TextIndexHandler.java @@ -24,13 +24,13 @@ import java.util.Map; import java.util.Set; import org.apache.pinot.segment.local.segment.index.dictionary.DictionaryIndexType; -import org.apache.pinot.segment.local.segment.index.forward.ForwardIndexType; import org.apache.pinot.segment.local.segment.index.loader.BaseIndexHandler; import org.apache.pinot.segment.local.segment.index.loader.SegmentPreProcessor; import org.apache.pinot.segment.spi.ColumnMetadata; import org.apache.pinot.segment.spi.creator.IndexCreationContext; import org.apache.pinot.segment.spi.index.FieldIndexConfigs; import org.apache.pinot.segment.spi.index.FieldIndexConfigsUtil; +import org.apache.pinot.segment.spi.index.IndexReaderFactory; import org.apache.pinot.segment.spi.index.StandardIndexes; import org.apache.pinot.segment.spi.index.TextIndexConfig; import org.apache.pinot.segment.spi.index.creator.TextIndexCreator; @@ -179,8 +179,9 @@ private void createTextIndexForColumn(SegmentDirectory.Writer segmentWriter, Col && _tableConfig.getIngestionConfig().isContinueOnError()) .build(); TextIndexConfig config = _fieldIndexConfigs.get(columnName).getConfig(StandardIndexes.text()); - - try (ForwardIndexReader forwardIndexReader = ForwardIndexType.read(segmentWriter, columnMetadata); + IndexReaderFactory readerFactory = StandardIndexes.forward().getReaderFactory(); + try (ForwardIndexReader forwardIndexReader = readerFactory.createIndexReader(segmentWriter, + _fieldIndexConfigs.get(columnMetadata.getColumnName()), columnMetadata); ForwardIndexReaderContext readerContext = forwardIndexReader.createContext(); TextIndexCreator textIndexCreator = StandardIndexes.text().createIndexCreator(context, config)) { if (columnMetadata.isSingleValue()) { diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/VectorIndexHandler.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/VectorIndexHandler.java index e614067ee51d..28a1b572bba1 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/VectorIndexHandler.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/loader/invertedindex/VectorIndexHandler.java @@ -24,13 +24,13 @@ import java.util.Map; import java.util.Set; import org.apache.commons.io.FileUtils; -import org.apache.pinot.segment.local.segment.index.forward.ForwardIndexType; import org.apache.pinot.segment.local.segment.index.loader.BaseIndexHandler; import org.apache.pinot.segment.spi.ColumnMetadata; import org.apache.pinot.segment.spi.V1Constants; import org.apache.pinot.segment.spi.creator.IndexCreationContext; import org.apache.pinot.segment.spi.index.FieldIndexConfigs; import org.apache.pinot.segment.spi.index.FieldIndexConfigsUtil; +import org.apache.pinot.segment.spi.index.IndexReaderFactory; import org.apache.pinot.segment.spi.index.StandardIndexes; import org.apache.pinot.segment.spi.index.creator.VectorIndexConfig; import org.apache.pinot.segment.spi.index.creator.VectorIndexCreator; @@ -208,7 +208,9 @@ private void handleNonDictionaryBasedColumn(SegmentDirectory.Writer segmentWrite && _tableConfig.getIngestionConfig().isContinueOnError()) .build(); VectorIndexConfig config = _fieldIndexConfigs.get(columnName).getConfig(StandardIndexes.vector()); - try (ForwardIndexReader forwardIndexReader = ForwardIndexType.read(segmentWriter, columnMetadata); + IndexReaderFactory readerFactory = StandardIndexes.forward().getReaderFactory(); + try (ForwardIndexReader forwardIndexReader = readerFactory.createIndexReader(segmentWriter, + _fieldIndexConfigs.get(columnMetadata.getColumnName()), columnMetadata); ForwardIndexReaderContext readerContext = forwardIndexReader.createContext(); VectorIndexCreator vectorIndexCreator = StandardIndexes.vector().createIndexCreator(context, config)) { int numDocs = columnMetadata.getTotalDocs(); diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/store/StarTreeLoaderUtils.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/store/StarTreeLoaderUtils.java index f7adc1a07d5c..06e511e65e21 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/store/StarTreeLoaderUtils.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/startree/v2/store/StarTreeLoaderUtils.java @@ -83,8 +83,8 @@ public static List loadStarTreeV2(SegmentDirectory.Reader segmentRea PinotDataBuffer forwardIndexDataBuffer = indexReader.getIndexFor(metric, StandardIndexes.forward()); DataType dataType = ValueAggregatorFactory.getAggregatedValueType(functionColumnPair.getFunctionType()); FieldSpec fieldSpec = new MetricFieldSpec(metric, dataType); - ForwardIndexReader forwardIndex = - ForwardIndexReaderFactory.createRawIndexReader(forwardIndexDataBuffer, dataType.getStoredType(), true); + ForwardIndexReader forwardIndex = ForwardIndexReaderFactory.getInstance() + .createRawIndexReader(forwardIndexDataBuffer, dataType.getStoredType(), true); dataSourceMap.put(metric, new StarTreeDataSource(fieldSpec, numDocs, forwardIndex, null)); } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/MultiValueFixedByteRawIndexCreatorTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/MultiValueFixedByteRawIndexCreatorTest.java index 1a876970d802..6e1b9d6e8197 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/MultiValueFixedByteRawIndexCreatorTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/MultiValueFixedByteRawIndexCreatorTest.java @@ -171,7 +171,8 @@ public void testMV(DataType dataType, List inputs, ToIntFunction sizeo } try (PinotDataBuffer buffer = PinotDataBuffer.mapFile(file, true, 0, file.length(), ByteOrder.BIG_ENDIAN, ""); - ForwardIndexReader reader = ForwardIndexReaderFactory.createRawIndexReader(buffer, dataType, false); + ForwardIndexReader reader = ForwardIndexReaderFactory.getInstance() + .createRawIndexReader(buffer, dataType, false); ForwardIndexReaderContext context = reader.createContext()) { T valueBuffer = constructor.apply(maxElements); for (int i = 0; i < numDocs; i++) { diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/MultiValueVarByteRawIndexCreatorTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/MultiValueVarByteRawIndexCreatorTest.java index 4f150b056722..96b973cdab26 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/MultiValueVarByteRawIndexCreatorTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/MultiValueVarByteRawIndexCreatorTest.java @@ -134,7 +134,8 @@ public void testMVString(ChunkCompressionType compressionType, boolean useFullSi } try (PinotDataBuffer buffer = PinotDataBuffer.mapFile(file, true, 0, file.length(), ByteOrder.BIG_ENDIAN, ""); - ForwardIndexReader reader = ForwardIndexReaderFactory.createRawIndexReader(buffer, DataType.STRING, false); + ForwardIndexReader reader = ForwardIndexReaderFactory.getInstance() + .createRawIndexReader(buffer, DataType.STRING, false); ForwardIndexReaderContext context = reader.createContext()) { String[] values = new String[maxElements]; for (int i = 0; i < numDocs; i++) { @@ -186,7 +187,8 @@ public void testMVBytes(ChunkCompressionType compressionType, boolean useFullSiz } try (PinotDataBuffer buffer = PinotDataBuffer.mapFile(file, true, 0, file.length(), ByteOrder.BIG_ENDIAN, ""); - ForwardIndexReader reader = ForwardIndexReaderFactory.createRawIndexReader(buffer, DataType.BYTES, false); + ForwardIndexReader reader = ForwardIndexReaderFactory.getInstance() + .createRawIndexReader(buffer, DataType.BYTES, false); ForwardIndexReaderContext context = reader.createContext()) { byte[][] values = new byte[maxElements][]; for (int i = 0; i < numDocs; i++) { diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/RawIndexCreatorTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/RawIndexCreatorTest.java index 848a62d29c28..85a0198b4615 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/RawIndexCreatorTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/creator/RawIndexCreatorTest.java @@ -169,8 +169,8 @@ public void testDoubleRawIndexCreator() public void testStringRawIndexCreator() throws Exception { PinotDataBuffer indexBuffer = getIndexBufferForColumn(STRING_COLUMN); - try ( - ForwardIndexReader rawIndexReader = ForwardIndexReaderFactory.createRawIndexReader(indexBuffer, DataType.STRING, + try (ForwardIndexReader rawIndexReader = ForwardIndexReaderFactory.getInstance() + .createRawIndexReader(indexBuffer, DataType.STRING, true); ForwardIndexReaderContext readerContext = rawIndexReader.createContext()) { _recordReader.rewind(); for (int row = 0; row < NUM_ROWS; row++) { @@ -207,8 +207,8 @@ private void testFixedLengthRawIndexCreator(String column, DataType dataType) public void testStringMVRawIndexCreator() throws Exception { PinotDataBuffer indexBuffer = getIndexBufferForColumn(STRING_MV_COLUMN); - try ( - ForwardIndexReader rawIndexReader = ForwardIndexReaderFactory.createRawIndexReader(indexBuffer, DataType.STRING, + try (ForwardIndexReader rawIndexReader = ForwardIndexReaderFactory.getInstance() + .createRawIndexReader(indexBuffer, DataType.STRING, false); ForwardIndexReaderContext readerContext = rawIndexReader.createContext()) { _recordReader.rewind(); int maxNumberOfMultiValues = @@ -240,7 +240,8 @@ public void testStringMVRawIndexCreator() public void testBytesMVRawIndexCreator() throws Exception { PinotDataBuffer indexBuffer = getIndexBufferForColumn(BYTES_MV_COLUMN); - try (ForwardIndexReader rawIndexReader = ForwardIndexReaderFactory.createRawIndexReader(indexBuffer, DataType.BYTES, + try (ForwardIndexReader rawIndexReader = ForwardIndexReaderFactory.getInstance() + .createRawIndexReader(indexBuffer, DataType.BYTES, false); ForwardIndexReaderContext readerContext = rawIndexReader.createContext()) { _recordReader.rewind(); int maxNumberOfMultiValues = diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/ForwardIndexHandlerTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/ForwardIndexHandlerTest.java index dd7825a67600..c7d7b733aa47 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/ForwardIndexHandlerTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/ForwardIndexHandlerTest.java @@ -39,7 +39,6 @@ import org.apache.commons.lang3.StringUtils; import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; import org.apache.pinot.segment.local.segment.index.dictionary.DictionaryIndexType; -import org.apache.pinot.segment.local.segment.index.forward.ForwardIndexType; import org.apache.pinot.segment.local.segment.index.loader.invertedindex.RangeIndexHandler; import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; import org.apache.pinot.segment.local.segment.readers.PinotSegmentColumnReader; @@ -49,7 +48,11 @@ import org.apache.pinot.segment.spi.compression.ChunkCompressionType; import org.apache.pinot.segment.spi.compression.DictIdCompressionType; import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.segment.spi.index.DictionaryIndexConfig; +import org.apache.pinot.segment.spi.index.FieldIndexConfigs; import org.apache.pinot.segment.spi.index.ForwardIndexConfig; +import org.apache.pinot.segment.spi.index.IndexReaderConstraintException; +import org.apache.pinot.segment.spi.index.IndexReaderFactory; import org.apache.pinot.segment.spi.index.IndexType; import org.apache.pinot.segment.spi.index.StandardIndexes; import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl; @@ -1120,8 +1123,11 @@ public void testChangeDictCompression() try (SegmentDirectory segmentDirectory = new SegmentLocalFSDirectory(INDEX_DIR, ReadMode.mmap); SegmentDirectory.Reader reader = segmentDirectory.createReader()) { - ForwardIndexReader forwardIndexReader = - ForwardIndexType.read(reader, segmentDirectory.getSegmentMetadata().getColumnMetadataFor(column)); + IndexReaderFactory readerFactory = StandardIndexes.forward().getReaderFactory(); + ColumnMetadata columnMetadata = segmentDirectory.getSegmentMetadata().getColumnMetadataFor(column); + FieldIndexConfigs fieldIndexConfigs = createFieldIndexConfigsFromMetadata(columnMetadata); + ForwardIndexReader forwardIndexReader = + readerFactory.createIndexReader(reader, fieldIndexConfigs, columnMetadata); assertTrue(forwardIndexReader.isDictionaryEncoded()); assertFalse(forwardIndexReader.isSingleValue()); assertEquals(forwardIndexReader.getDictIdCompressionType(), DictIdCompressionType.MV_ENTRY_DICT); @@ -1147,8 +1153,11 @@ public void testChangeDictCompression() try (SegmentDirectory segmentDirectory = new SegmentLocalFSDirectory(INDEX_DIR, ReadMode.mmap); SegmentDirectory.Reader reader = segmentDirectory.createReader()) { - ForwardIndexReader forwardIndexReader = - ForwardIndexType.read(reader, segmentDirectory.getSegmentMetadata().getColumnMetadataFor(column)); + IndexReaderFactory readerFactory = StandardIndexes.forward().getReaderFactory(); + ColumnMetadata columnMetadata = segmentDirectory.getSegmentMetadata().getColumnMetadataFor(column); + FieldIndexConfigs fieldIndexConfigs = createFieldIndexConfigsFromMetadata(columnMetadata); + ForwardIndexReader forwardIndexReader = + readerFactory.createIndexReader(reader, fieldIndexConfigs, columnMetadata); assertTrue(forwardIndexReader.isDictionaryEncoded()); assertFalse(forwardIndexReader.isSingleValue()); assertNull(forwardIndexReader.getDictIdCompressionType()); @@ -2019,12 +2028,14 @@ private void validateForwardIndex(String columnName, @Nullable CompressionCodec try (SegmentDirectory segmentDirectory = new SegmentLocalFSDirectory(INDEX_DIR, ReadMode.mmap); SegmentDirectory.Reader reader = segmentDirectory.createReader()) { validateForwardIndex(segmentDirectory, reader, columnName, expectedCompressionType, isSorted); + } catch (IndexReaderConstraintException e) { + throw new RuntimeException(e); } } private void validateForwardIndex(SegmentDirectory segmentDirectory, SegmentDirectory.Reader reader, String columnName, @Nullable CompressionCodec expectedCompressionType, boolean isSorted) - throws IOException { + throws IOException, IndexReaderConstraintException { ColumnMetadata columnMetadata = segmentDirectory.getSegmentMetadata().getColumnMetadataFor(columnName); boolean isSingleValue = columnMetadata.isSingleValue(); @@ -2036,7 +2047,9 @@ private void validateForwardIndex(SegmentDirectory segmentDirectory, SegmentDire assertTrue(reader.hasIndexFor(columnName, StandardIndexes.forward())); // Check Compression type in header - ForwardIndexReader fwdIndexReader = ForwardIndexType.read(reader, columnMetadata); + IndexReaderFactory readerFactory = StandardIndexes.forward().getReaderFactory(); + FieldIndexConfigs fieldIndexConfigs = createFieldIndexConfigsFromMetadata(columnMetadata); + ForwardIndexReader fwdIndexReader = readerFactory.createIndexReader(reader, fieldIndexConfigs, columnMetadata); ChunkCompressionType fwdIndexCompressionType = fwdIndexReader.getCompressionType(); if (expectedCompressionType != null) { assertNotNull(fwdIndexCompressionType); @@ -2044,8 +2057,9 @@ private void validateForwardIndex(SegmentDirectory segmentDirectory, SegmentDire } else { assertNull(fwdIndexCompressionType); } - - try (ForwardIndexReader forwardIndexReader = ForwardIndexType.read(reader, columnMetadata)) { + fieldIndexConfigs = createFieldIndexConfigsFromMetadata(columnMetadata); + try (ForwardIndexReader forwardIndexReader = readerFactory.createIndexReader(reader, fieldIndexConfigs, + columnMetadata)) { Dictionary dictionary = null; if (columnMetadata.hasDictionary()) { dictionary = DictionaryIndexType.read(reader, columnMetadata); @@ -2257,4 +2271,20 @@ private void validateMetadataProperties(String column, boolean hasDictionary, in assertEquals(columnMetadata.getMinValue(), minValue); assertEquals(columnMetadata.getMaxValue(), maxValue); } + + private FieldIndexConfigs createFieldIndexConfigsFromMetadata(ColumnMetadata columnMetadata) { + FieldIndexConfigs.Builder builder = new FieldIndexConfigs.Builder(); + + // Add forward index config + ForwardIndexConfig forwardIndexConfig = ForwardIndexConfig.getDefault(); + builder.add(StandardIndexes.forward(), forwardIndexConfig); + + // Add dictionary config if the column has dictionary + if (columnMetadata.hasDictionary()) { + DictionaryIndexConfig dictionaryConfig = DictionaryIndexConfig.DEFAULT; + builder.add(StandardIndexes.dictionary(), dictionaryConfig); + } + + return builder.build(); + } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/SegmentPreProcessorTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/SegmentPreProcessorTest.java index 53b2564a68b7..f53f3ba35698 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/SegmentPreProcessorTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/segment/index/loader/SegmentPreProcessorTest.java @@ -39,7 +39,6 @@ import org.apache.pinot.segment.local.segment.creator.SegmentTestUtils; import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; import org.apache.pinot.segment.local.segment.index.converter.SegmentV1V2ToV3FormatConverter; -import org.apache.pinot.segment.local.segment.index.forward.ForwardIndexType; import org.apache.pinot.segment.local.segment.index.loader.columnminmaxvalue.ColumnMinMaxValueGeneratorMode; import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; import org.apache.pinot.segment.local.segment.store.SegmentLocalFSDirectory; @@ -53,6 +52,9 @@ import org.apache.pinot.segment.spi.compression.ChunkCompressionType; import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; import org.apache.pinot.segment.spi.creator.SegmentVersion; +import org.apache.pinot.segment.spi.index.FieldIndexConfigs; +import org.apache.pinot.segment.spi.index.ForwardIndexConfig; +import org.apache.pinot.segment.spi.index.IndexReaderFactory; import org.apache.pinot.segment.spi.index.IndexType; import org.apache.pinot.segment.spi.index.StandardIndexes; import org.apache.pinot.segment.spi.index.metadata.SegmentMetadataImpl; @@ -805,8 +807,12 @@ private void validateIndex(IndexType indexType, String column, int card // Check if the raw forward index compressionType is correct. if (expectedCompressionType != null) { assertFalse(hasDictionary); - - try (ForwardIndexReader fwdIndexReader = ForwardIndexType.read(reader, columnMetadata)) { + IndexReaderFactory readerFactory = StandardIndexes.forward().getReaderFactory(); + FieldIndexConfigs fieldIndexConfigs = new FieldIndexConfigs.Builder() + .add(StandardIndexes.forward(), ForwardIndexConfig.getDefault()) + .build(); + try (ForwardIndexReader fwdIndexReader = readerFactory.createIndexReader(reader, fieldIndexConfigs, + columnMetadata)) { ChunkCompressionType compressionType = fwdIndexReader.getCompressionType(); assertEquals(compressionType, expectedCompressionType); } From 30119129f8376b376df1d5dc8572cf41459e97d6 Mon Sep 17 00:00:00 2001 From: Hongkun Xu Date: Thu, 24 Jul 2025 20:26:59 +0800 Subject: [PATCH 016/167] Support executing backfill jobs on clusters that use self-signed certificates (#16411) Signed-off-by: Hongkun Xu --- .../generation/SegmentGenerationUtils.java | 68 +++++++++++++++---- .../SparkSegmentGenerationJobRunner.java | 6 +- .../SparkSegmentGenerationJobRunner.java | 6 +- .../spi/ingestion/batch/spec/TlsSpec.java | 18 +++++ 4 files changed, 82 insertions(+), 16 deletions(-) diff --git a/pinot-common/src/main/java/org/apache/pinot/common/segment/generation/SegmentGenerationUtils.java b/pinot-common/src/main/java/org/apache/pinot/common/segment/generation/SegmentGenerationUtils.java index 981bd4d464d8..b40eff46ebae 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/segment/generation/SegmentGenerationUtils.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/segment/generation/SegmentGenerationUtils.java @@ -23,6 +23,7 @@ import java.io.File; import java.io.IOException; import java.io.InputStream; +import java.net.HttpURLConnection; import java.net.URI; import java.net.URISyntaxException; import java.net.URL; @@ -31,15 +32,21 @@ import java.nio.file.FileSystems; import java.nio.file.PathMatcher; import java.nio.file.Paths; +import java.security.SecureRandom; import java.util.ArrayList; import java.util.List; import javax.annotation.Nullable; +import javax.net.ssl.HttpsURLConnection; +import javax.net.ssl.SSLContext; +import javax.net.ssl.TrustManagerFactory; import org.apache.commons.io.IOUtils; import org.apache.commons.lang3.StringUtils; +import org.apache.pinot.common.utils.tls.TlsUtils; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.data.Schema; import org.apache.pinot.spi.filesystem.PinotFS; import org.apache.pinot.spi.filesystem.PinotFSFactory; +import org.apache.pinot.spi.ingestion.batch.spec.TlsSpec; import org.apache.pinot.spi.utils.JsonUtils; @@ -65,6 +72,10 @@ public static Schema getSchema(String schemaURIString) { } public static Schema getSchema(String schemaURIString, String authToken) { + return getSchema(schemaURIString, authToken, null); + } + + public static Schema getSchema(String schemaURIString, String authToken, TlsSpec tlsSpec) { URI schemaURI; try { schemaURIString = sanitizeURIString(schemaURIString); @@ -90,7 +101,7 @@ public static Schema getSchema(String schemaURIString, String authToken) { } else { // Try to directly read from URI. try { - schemaJson = fetchUrl(schemaURI.toURL(), authToken); + schemaJson = fetchUrl(schemaURI.toURL(), authToken, tlsSpec); } catch (IOException e) { throw new RuntimeException("Failed to read from Schema URI - '" + schemaURI + "'", e); } @@ -108,6 +119,10 @@ public static TableConfig getTableConfig(String tableConfigURIStr) { } public static TableConfig getTableConfig(String tableConfigURIStr, String authToken) { + return getTableConfig(tableConfigURIStr, authToken, null); + } + + public static TableConfig getTableConfig(String tableConfigURIStr, String authToken, TlsSpec tlsSpec) { URI tableConfigURI; try { tableConfigURI = new URI(tableConfigURIStr); @@ -125,7 +140,7 @@ public static TableConfig getTableConfig(String tableConfigURIStr, String authTo } } else { try { - tableConfigJson = fetchUrl(tableConfigURI.toURL(), authToken); + tableConfigJson = fetchUrl(tableConfigURI.toURL(), authToken, tlsSpec); } catch (IOException e) { throw new RuntimeException( "Failed to read from table config file data stream on Pinot fs - '" + tableConfigURI + "'", e); @@ -225,21 +240,50 @@ public static URI getDirectoryURI(String uriStr) } /** - * Retrieve a URL via GET request, with an optional authorization token. + * Retrieves the content of the given URL using a GET request. + * Supports HTTPS connections, including those with self-signed certificates. * - * @param url target url - * @param authToken optional auth token, or null - * @return fetched document - * @throws IOException on connection problems + * @param url The target URL to fetch. + * @param authToken Optional authorization token to include in the request header, or null. + * @return The response body as a string. + * @throws IOException If an error occurs during the connection or reading the response. */ - private static String fetchUrl(URL url, String authToken) + public static String fetchUrl(URL url, String authToken, TlsSpec tlsSpec) throws IOException { - URLConnection connection = url.openConnection(); + try { + URLConnection connection = url.openConnection(); + + if (connection instanceof HttpsURLConnection) { + HttpsURLConnection httpsConn = (HttpsURLConnection) connection; + + if (tlsSpec != null) { + TrustManagerFactory tmf = TlsUtils.createTrustManagerFactory( + tlsSpec.getTrustStorePath(), + tlsSpec.getTrustStorePassword(), + tlsSpec.getTrustStoreType()); + + SSLContext sslContext = SSLContext.getInstance("TLS"); + sslContext.init(null, tmf.getTrustManagers(), new SecureRandom()); + + httpsConn.setSSLSocketFactory(sslContext.getSocketFactory()); + httpsConn.setConnectTimeout(tlsSpec.getConnectTimeout()); + httpsConn.setReadTimeout(tlsSpec.getReadTimeout()); + } + connection = httpsConn; + } + + if (StringUtils.isNotBlank(authToken)) { + connection.setRequestProperty("Authorization", authToken); + } + + if (connection instanceof HttpURLConnection) { + ((HttpURLConnection) connection).setRequestMethod("GET"); + } - if (StringUtils.isNotBlank(authToken)) { - connection.setRequestProperty("Authorization", authToken); + return IOUtils.toString(connection.getInputStream(), StandardCharsets.UTF_8); + } catch (Exception e) { + throw new IOException("Failed to fetch URL: " + url, e); } - return IOUtils.toString(connection.getInputStream(), StandardCharsets.UTF_8); } diff --git a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-2.4/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark/SparkSegmentGenerationJobRunner.java b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-2.4/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark/SparkSegmentGenerationJobRunner.java index 73581206d788..bc3a42e4dff7 100644 --- a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-2.4/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark/SparkSegmentGenerationJobRunner.java +++ b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-2.4/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark/SparkSegmentGenerationJobRunner.java @@ -269,9 +269,11 @@ public void call(String pathAndIdx) taskSpec.setOutputDirectoryPath(localOutputTempDir.getAbsolutePath()); taskSpec.setRecordReaderSpec(_spec.getRecordReaderSpec()); taskSpec - .setSchema(SegmentGenerationUtils.getSchema(_spec.getTableSpec().getSchemaURI(), _spec.getAuthToken())); + .setSchema(SegmentGenerationUtils.getSchema(_spec.getTableSpec().getSchemaURI(), _spec.getAuthToken(), + _spec.getTlsSpec())); taskSpec.setTableConfig( - SegmentGenerationUtils.getTableConfig(_spec.getTableSpec().getTableConfigURI(), _spec.getAuthToken())); + SegmentGenerationUtils.getTableConfig(_spec.getTableSpec().getTableConfigURI(), _spec.getAuthToken(), + _spec.getTlsSpec())); taskSpec.setSequenceId(idx); taskSpec.setSegmentNameGeneratorSpec(_spec.getSegmentNameGeneratorSpec()); taskSpec.setFailOnEmptySegment(_spec.isFailOnEmptySegment()); diff --git a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-3/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark3/SparkSegmentGenerationJobRunner.java b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-3/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark3/SparkSegmentGenerationJobRunner.java index a40bbf652e36..ed1f4d1fe5d4 100644 --- a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-3/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark3/SparkSegmentGenerationJobRunner.java +++ b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-3/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark3/SparkSegmentGenerationJobRunner.java @@ -278,9 +278,11 @@ public void call(String pathAndIdx) taskSpec.setOutputDirectoryPath(localOutputTempDir.getAbsolutePath()); taskSpec.setRecordReaderSpec(_spec.getRecordReaderSpec()); taskSpec.setSchema( - SegmentGenerationUtils.getSchema(_spec.getTableSpec().getSchemaURI(), _spec.getAuthToken())); + SegmentGenerationUtils.getSchema(_spec.getTableSpec().getSchemaURI(), _spec.getAuthToken(), + _spec.getTlsSpec())); taskSpec.setTableConfig( - SegmentGenerationUtils.getTableConfig(_spec.getTableSpec().getTableConfigURI(), _spec.getAuthToken())); + SegmentGenerationUtils.getTableConfig(_spec.getTableSpec().getTableConfigURI(), _spec.getAuthToken(), + _spec.getTlsSpec())); taskSpec.setSequenceId(idx); taskSpec.setSegmentNameGeneratorSpec(_spec.getSegmentNameGeneratorSpec()); taskSpec.setFailOnEmptySegment(_spec.isFailOnEmptySegment()); diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/ingestion/batch/spec/TlsSpec.java b/pinot-spi/src/main/java/org/apache/pinot/spi/ingestion/batch/spec/TlsSpec.java index bcc1f890206c..e2568b218779 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/ingestion/batch/spec/TlsSpec.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/ingestion/batch/spec/TlsSpec.java @@ -32,6 +32,8 @@ public class TlsSpec implements Serializable { private String _trustStorePath; private String _trustStorePassword; private String _keyStoreType; + private int _connectTimeout = 5000; + private int _readTimeout = 5000; public String getKeyStorePath() { return _keyStorePath; @@ -80,4 +82,20 @@ public String getKeyStoreType() { public void setKeyStoreType(String keyStoreType) { _keyStoreType = keyStoreType; } + + public int getConnectTimeout() { + return _connectTimeout; + } + + public void setConnectTimeout(int connectTimeout) { + _connectTimeout = connectTimeout; + } + + public int getReadTimeout() { + return _readTimeout; + } + + public void setReadTimeout(int readTimeout) { + _readTimeout = readTimeout; + } } From 9a18becd1591340f653b1e6caf53ba069b1645e4 Mon Sep 17 00:00:00 2001 From: Rajat Venkatesh <1638298+vrajat@users.noreply.github.com> Date: Thu, 24 Jul 2025 18:35:40 +0530 Subject: [PATCH 017/167] Use correlation ID instead of request id in PerQueryCpuMemAccountant (#16040) --- .../broker/helix/BaseBrokerStarter.java | 33 ++++++++++--------- .../BaseBrokerRequestHandler.java | 7 +++- .../BaseSingleStageBrokerRequestHandler.java | 9 +++-- .../GrpcBrokerRequestHandler.java | 5 +-- .../MultiStageBrokerRequestHandler.java | 10 +++--- .../SingleConnectionBrokerRequestHandler.java | 8 +++-- .../TimeSeriesRequestHandler.java | 5 +-- ...seSingleStageBrokerRequestHandlerTest.java | 4 ++- .../LiteralOnlyBrokerRequestTest.java | 10 ++++-- .../query/reduce/BrokerReduceService.java | 12 ++++++- .../query/reduce/GroupByDataTableReducer.java | 7 ++-- .../query/reduce/ResultReducerFactory.java | 5 +-- .../core/query/scheduler/QueryScheduler.java | 2 +- .../OfflineGRPCServerIntegrationTest.java | 4 ++- .../query/service/server/QueryServer.java | 4 +-- 15 files changed, 81 insertions(+), 44 deletions(-) diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java index 3ae7fd846ce4..b127e61a6d54 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java @@ -347,6 +347,19 @@ public void start() instanceId -> _routingManager.excludeServerFromRouting(instanceId)); _failureDetector.start(); + // Enable/disable thread CPU time measurement through instance config. + ThreadResourceUsageProvider.setThreadCpuTimeMeasurementEnabled( + _brokerConf.getProperty(CommonConstants.Broker.CONFIG_OF_ENABLE_THREAD_CPU_TIME_MEASUREMENT, + CommonConstants.Broker.DEFAULT_ENABLE_THREAD_CPU_TIME_MEASUREMENT)); + // Enable/disable thread memory allocation tracking through instance config + ThreadResourceUsageProvider.setThreadMemoryMeasurementEnabled( + _brokerConf.getProperty(CommonConstants.Broker.CONFIG_OF_ENABLE_THREAD_ALLOCATED_BYTES_MEASUREMENT, + CommonConstants.Broker.DEFAULT_THREAD_ALLOCATED_BYTES_MEASUREMENT)); + _resourceUsageAccountant = Tracing.ThreadAccountantOps.createThreadAccountant( + _brokerConf.subset(CommonConstants.PINOT_QUERY_SCHEDULER_PREFIX), _instanceId, + org.apache.pinot.spi.config.instance.InstanceType.BROKER); + Preconditions.checkNotNull(_resourceUsageAccountant); + // Create Broker request handler. String brokerId = _brokerConf.getProperty(Broker.CONFIG_OF_BROKER_ID, getDefaultBrokerId()); String brokerRequestHandlerType = @@ -354,7 +367,7 @@ public void start() BaseSingleStageBrokerRequestHandler singleStageBrokerRequestHandler; if (brokerRequestHandlerType.equalsIgnoreCase(Broker.GRPC_BROKER_REQUEST_HANDLER_TYPE)) { singleStageBrokerRequestHandler = new GrpcBrokerRequestHandler(_brokerConf, brokerId, _routingManager, - _accessControlFactory, _queryQuotaManager, tableCache, _failureDetector); + _accessControlFactory, _queryQuotaManager, tableCache, _failureDetector, _resourceUsageAccountant); } else { // Default request handler type, i.e. netty NettyConfig nettyDefaults = NettyConfig.extractNettyConfig(_brokerConf, Broker.BROKER_NETTY_PREFIX); @@ -366,7 +379,7 @@ public void start() singleStageBrokerRequestHandler = new SingleConnectionBrokerRequestHandler(_brokerConf, brokerId, _routingManager, _accessControlFactory, _queryQuotaManager, tableCache, nettyDefaults, tlsDefaults, _serverRoutingStatsManager, - _failureDetector); + _failureDetector, _resourceUsageAccountant); } MultiStageBrokerRequestHandler multiStageBrokerRequestHandler = null; QueryDispatcher queryDispatcher = null; @@ -379,13 +392,13 @@ public void start() queryDispatcher = createQueryDispatcher(_brokerConf); multiStageBrokerRequestHandler = new MultiStageBrokerRequestHandler(_brokerConf, brokerId, _routingManager, _accessControlFactory, - _queryQuotaManager, tableCache, _multiStageQueryThrottler, _failureDetector); + _queryQuotaManager, tableCache, _multiStageQueryThrottler, _failureDetector, _resourceUsageAccountant); } TimeSeriesRequestHandler timeSeriesRequestHandler = null; if (StringUtils.isNotBlank(_brokerConf.getProperty(PinotTimeSeriesConfiguration.getEnabledLanguagesConfigKey()))) { Preconditions.checkNotNull(queryDispatcher, "Multistage Engine should be enabled to use time-series engine"); timeSeriesRequestHandler = new TimeSeriesRequestHandler(_brokerConf, brokerId, _routingManager, - _accessControlFactory, _queryQuotaManager, tableCache, queryDispatcher); + _accessControlFactory, _queryQuotaManager, tableCache, queryDispatcher, _resourceUsageAccountant); } LOGGER.info("Initializing PinotFSFactory"); @@ -409,18 +422,6 @@ public void start() timeSeriesRequestHandler, _responseStore); _brokerRequestHandler.start(); - // Enable/disable thread CPU time measurement through instance config. - ThreadResourceUsageProvider.setThreadCpuTimeMeasurementEnabled( - _brokerConf.getProperty(CommonConstants.Broker.CONFIG_OF_ENABLE_THREAD_CPU_TIME_MEASUREMENT, - CommonConstants.Broker.DEFAULT_ENABLE_THREAD_CPU_TIME_MEASUREMENT)); - // Enable/disable thread memory allocation tracking through instance config - ThreadResourceUsageProvider.setThreadMemoryMeasurementEnabled( - _brokerConf.getProperty(CommonConstants.Broker.CONFIG_OF_ENABLE_THREAD_ALLOCATED_BYTES_MEASUREMENT, - CommonConstants.Broker.DEFAULT_THREAD_ALLOCATED_BYTES_MEASUREMENT)); - _resourceUsageAccountant = Tracing.ThreadAccountantOps.createThreadAccountant( - _brokerConf.subset(CommonConstants.PINOT_QUERY_SCHEDULER_PREFIX), _instanceId, - org.apache.pinot.spi.config.instance.InstanceType.BROKER); - Preconditions.checkNotNull(_resourceUsageAccountant); Tracing.ThreadAccountantOps.startThreadAccountant(); String controllerUrl = _brokerConf.getProperty(Broker.CONTROLLER_URL); diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java index 895102dc61ea..b4aafbea3f23 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java @@ -53,6 +53,7 @@ import org.apache.pinot.common.utils.request.RequestUtils; import org.apache.pinot.core.auth.Actions; import org.apache.pinot.core.auth.TargetType; +import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant; import org.apache.pinot.spi.auth.AuthorizationResult; import org.apache.pinot.spi.auth.TableAuthorizationResult; import org.apache.pinot.spi.auth.broker.RequesterIdentity; @@ -89,6 +90,8 @@ public abstract class BaseBrokerRequestHandler implements BrokerRequestHandler { protected final QueryLogger _queryLogger; @Nullable protected final String _enableNullHandling; + protected final ThreadResourceUsageAccountant _resourceUsageAccountant; + /** * Maps broker-generated query id to the query string. */ @@ -99,7 +102,8 @@ public abstract class BaseBrokerRequestHandler implements BrokerRequestHandler { protected final Map _clientQueryIds; public BaseBrokerRequestHandler(PinotConfiguration config, String brokerId, BrokerRoutingManager routingManager, - AccessControlFactory accessControlFactory, QueryQuotaManager queryQuotaManager, TableCache tableCache) { + AccessControlFactory accessControlFactory, QueryQuotaManager queryQuotaManager, TableCache tableCache, + ThreadResourceUsageAccountant resourceUsageAccountant) { _config = config; _brokerId = brokerId; _routingManager = routingManager; @@ -125,6 +129,7 @@ public BaseBrokerRequestHandler(PinotConfiguration config, String brokerId, Brok _queriesById = null; _clientQueryIds = null; } + _resourceUsageAccountant = resourceUsageAccountant; } @Override diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java index 5c6fefe40b2f..d9a1e365ead6 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java @@ -99,6 +99,8 @@ import org.apache.pinot.query.routing.table.LogicalTableRouteProvider; import org.apache.pinot.query.routing.table.TableRouteProvider; import org.apache.pinot.segment.local.function.GroovyFunctionEvaluator; +import org.apache.pinot.spi.accounting.ThreadExecutionContext; +import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant; import org.apache.pinot.spi.auth.AuthorizationResult; import org.apache.pinot.spi.auth.TableRowColAccessResult; import org.apache.pinot.spi.auth.broker.RequesterIdentity; @@ -169,8 +171,8 @@ public abstract class BaseSingleStageBrokerRequestHandler extends BaseBrokerRequ public BaseSingleStageBrokerRequestHandler(PinotConfiguration config, String brokerId, BrokerRoutingManager routingManager, AccessControlFactory accessControlFactory, - QueryQuotaManager queryQuotaManager, TableCache tableCache) { - super(config, brokerId, routingManager, accessControlFactory, queryQuotaManager, tableCache); + QueryQuotaManager queryQuotaManager, TableCache tableCache, ThreadResourceUsageAccountant accountant) { + super(config, brokerId, routingManager, accessControlFactory, queryQuotaManager, tableCache, accountant); _disableGroovy = _config.getProperty(Broker.DISABLE_GROOVY, Broker.DEFAULT_DISABLE_GROOVY); _useApproximateFunction = _config.getProperty(Broker.USE_APPROXIMATE_FUNCTION, false); _defaultHllLog2m = _config.getProperty(CommonConstants.Helix.DEFAULT_HYPERLOGLOG_LOG2M_KEY, @@ -323,7 +325,8 @@ protected BrokerResponse handleRequest(long requestId, String query, SqlNodeAndO //Start instrumentation context. This must not be moved further below interspersed into the code. String workloadName = QueryOptionsUtils.getWorkloadName(sqlNodeAndOptions.getOptions()); - Tracing.ThreadAccountantOps.setupRunner(String.valueOf(requestId), workloadName); + _resourceUsageAccountant.setupRunner(QueryThreadContext.getCid(), CommonConstants.Accounting.ANCHOR_TASK_ID, + ThreadExecutionContext.TaskType.SSE, workloadName); try { return doHandleRequest(requestId, query, sqlNodeAndOptions, request, requesterIdentity, requestContext, diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java index 66973a842c64..0c832085ee26 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java @@ -42,6 +42,7 @@ import org.apache.pinot.core.transport.ServerInstance; import org.apache.pinot.core.transport.ServerRoutingInstance; import org.apache.pinot.core.transport.TableRouteInfo; +import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.env.PinotConfiguration; import org.apache.pinot.spi.query.QueryThreadContext; @@ -64,8 +65,8 @@ public class GrpcBrokerRequestHandler extends BaseSingleStageBrokerRequestHandle // TODO: Support TLS public GrpcBrokerRequestHandler(PinotConfiguration config, String brokerId, BrokerRoutingManager routingManager, AccessControlFactory accessControlFactory, QueryQuotaManager queryQuotaManager, TableCache tableCache, - FailureDetector failureDetector) { - super(config, brokerId, routingManager, accessControlFactory, queryQuotaManager, tableCache); + FailureDetector failureDetector, ThreadResourceUsageAccountant accountant) { + super(config, brokerId, routingManager, accessControlFactory, queryQuotaManager, tableCache, accountant); _streamingReduceService = new StreamingReduceService(config); _streamingQueryClient = new PinotServerStreamingQueryClient(GrpcConfig.buildGrpcQueryConfig(config)); _failureDetector = failureDetector; diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java index 1a4563bddd2d..1fc60da9d286 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java @@ -85,6 +85,7 @@ import org.apache.pinot.query.runtime.plan.MultiStageQueryStats; import org.apache.pinot.query.service.dispatch.QueryDispatcher; import org.apache.pinot.spi.accounting.ThreadExecutionContext; +import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant; import org.apache.pinot.spi.auth.TableAuthorizationResult; import org.apache.pinot.spi.auth.broker.RequesterIdentity; import org.apache.pinot.spi.env.PinotConfiguration; @@ -133,8 +134,9 @@ public class MultiStageBrokerRequestHandler extends BaseBrokerRequestHandler { public MultiStageBrokerRequestHandler(PinotConfiguration config, String brokerId, BrokerRoutingManager routingManager, AccessControlFactory accessControlFactory, QueryQuotaManager queryQuotaManager, TableCache tableCache, - MultiStageQueryThrottler queryThrottler, FailureDetector failureDetector) { - super(config, brokerId, routingManager, accessControlFactory, queryQuotaManager, tableCache); + MultiStageQueryThrottler queryThrottler, FailureDetector failureDetector, + ThreadResourceUsageAccountant accountant) { + super(config, brokerId, routingManager, accessControlFactory, queryQuotaManager, tableCache, accountant); String hostname = config.getProperty(CommonConstants.MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_HOSTNAME); int port = Integer.parseInt(config.getProperty(CommonConstants.MultiStageQueryRunner.KEY_OF_QUERY_RUNNER_PORT)); _workerManager = new WorkerManager(_brokerId, hostname, port, _routingManager); @@ -528,8 +530,8 @@ private BrokerResponse query(QueryEnvironment.CompiledQuery query, long requestI try { String workloadName = QueryOptionsUtils.getWorkloadName(query.getOptions()); - Tracing.ThreadAccountantOps.setupRunner(String.valueOf(requestId), ThreadExecutionContext.TaskType.MSE, - workloadName); + _resourceUsageAccountant.setupRunner(QueryThreadContext.getCid(), CommonConstants.Accounting.ANCHOR_TASK_ID, + ThreadExecutionContext.TaskType.MSE, workloadName); long executionStartTimeNs = System.nanoTime(); QueryDispatcher.QueryResult queryResults; diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java index e870a96a5ec7..e636a2ce5b7d 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java @@ -47,6 +47,7 @@ import org.apache.pinot.core.transport.ServerRoutingInstance; import org.apache.pinot.core.transport.TableRouteInfo; import org.apache.pinot.core.transport.server.routing.stats.ServerRoutingStatsManager; +import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant; import org.apache.pinot.spi.env.PinotConfiguration; import org.apache.pinot.spi.exception.QueryErrorCode; import org.apache.pinot.spi.trace.RequestContext; @@ -71,9 +72,10 @@ public class SingleConnectionBrokerRequestHandler extends BaseSingleStageBrokerR public SingleConnectionBrokerRequestHandler(PinotConfiguration config, String brokerId, BrokerRoutingManager routingManager, AccessControlFactory accessControlFactory, QueryQuotaManager queryQuotaManager, TableCache tableCache, NettyConfig nettyConfig, TlsConfig tlsConfig, - ServerRoutingStatsManager serverRoutingStatsManager, FailureDetector failureDetector) { - super(config, brokerId, routingManager, accessControlFactory, queryQuotaManager, tableCache); - _brokerReduceService = new BrokerReduceService(_config); + ServerRoutingStatsManager serverRoutingStatsManager, FailureDetector failureDetector, + ThreadResourceUsageAccountant accountant) { + super(config, brokerId, routingManager, accessControlFactory, queryQuotaManager, tableCache, accountant); + _brokerReduceService = new BrokerReduceService(_config, accountant); _queryRouter = new QueryRouter(_brokerId, _brokerMetrics, nettyConfig, tlsConfig, serverRoutingStatsManager); _failureDetector = failureDetector; _failureDetector.registerUnhealthyServerRetrier(this::retryUnhealthyServer); diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/TimeSeriesRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/TimeSeriesRequestHandler.java index 76d382a41b4b..fb38c572cf5f 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/TimeSeriesRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/TimeSeriesRequestHandler.java @@ -49,6 +49,7 @@ import org.apache.pinot.core.auth.Actions; import org.apache.pinot.core.auth.TargetType; import org.apache.pinot.query.service.dispatch.QueryDispatcher; +import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant; import org.apache.pinot.spi.auth.AuthorizationResult; import org.apache.pinot.spi.auth.broker.RequesterIdentity; import org.apache.pinot.spi.env.PinotConfiguration; @@ -71,8 +72,8 @@ public class TimeSeriesRequestHandler extends BaseBrokerRequestHandler { public TimeSeriesRequestHandler(PinotConfiguration config, String brokerId, BrokerRoutingManager routingManager, AccessControlFactory accessControlFactory, QueryQuotaManager queryQuotaManager, TableCache tableCache, - QueryDispatcher queryDispatcher) { - super(config, brokerId, routingManager, accessControlFactory, queryQuotaManager, tableCache); + QueryDispatcher queryDispatcher, ThreadResourceUsageAccountant accountant) { + super(config, brokerId, routingManager, accessControlFactory, queryQuotaManager, tableCache, accountant); _queryEnvironment = new TimeSeriesQueryEnvironment(config, routingManager, tableCache); _queryEnvironment.init(config); _queryDispatcher = queryDispatcher; diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandlerTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandlerTest.java index d408171c92c2..a3e80e0e882e 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandlerTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandlerTest.java @@ -44,6 +44,7 @@ import org.apache.pinot.spi.eventlistener.query.BrokerQueryEventListenerFactory; import org.apache.pinot.spi.exception.BadQueryRequestException; import org.apache.pinot.spi.trace.RequestContext; +import org.apache.pinot.spi.trace.Tracing; import org.apache.pinot.spi.utils.CommonConstants.Broker; import org.apache.pinot.sql.parsers.CalciteSqlParser; import org.apache.pinot.util.TestUtils; @@ -182,7 +183,8 @@ public void testCancelQuery() { BrokerQueryEventListenerFactory.init(config); BaseSingleStageBrokerRequestHandler requestHandler = new BaseSingleStageBrokerRequestHandler(config, "testBrokerId", routingManager, - new AllowAllAccessControlFactory(), queryQuotaManager, tableCache) { + new AllowAllAccessControlFactory(), queryQuotaManager, tableCache, + new Tracing.DefaultThreadResourceUsageAccountant()) { @Override public void start() { } diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/LiteralOnlyBrokerRequestTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/LiteralOnlyBrokerRequestTest.java index a6c1a7c16443..9614f29c8efa 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/LiteralOnlyBrokerRequestTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/LiteralOnlyBrokerRequestTest.java @@ -32,6 +32,7 @@ import org.apache.pinot.spi.env.PinotConfiguration; import org.apache.pinot.spi.eventlistener.query.BrokerQueryEventListenerFactory; import org.apache.pinot.spi.exception.QueryErrorCode; +import org.apache.pinot.spi.trace.Tracing; import org.apache.pinot.spi.utils.BytesUtils; import org.apache.pinot.sql.parsers.CalciteSqlParser; import org.testng.annotations.BeforeClass; @@ -171,7 +172,8 @@ public void testBrokerRequestHandler() throws Exception { SingleConnectionBrokerRequestHandler requestHandler = new SingleConnectionBrokerRequestHandler(new PinotConfiguration(), "testBrokerId", null, ACCESS_CONTROL_FACTORY, - null, null, null, null, mock(ServerRoutingStatsManager.class), mock(FailureDetector.class)); + null, null, null, null, mock(ServerRoutingStatsManager.class), mock(FailureDetector.class), + new Tracing.DefaultThreadResourceUsageAccountant()); long randNum = RANDOM.nextLong(); byte[] randBytes = new byte[12]; @@ -195,7 +197,8 @@ public void testBrokerRequestHandlerWithAsFunction() throws Exception { SingleConnectionBrokerRequestHandler requestHandler = new SingleConnectionBrokerRequestHandler(new PinotConfiguration(), "testBrokerId", null, ACCESS_CONTROL_FACTORY, - null, null, null, null, mock(ServerRoutingStatsManager.class), mock(FailureDetector.class)); + null, null, null, null, mock(ServerRoutingStatsManager.class), mock(FailureDetector.class), + new Tracing.DefaultThreadResourceUsageAccountant()); long currentTsMin = System.currentTimeMillis(); BrokerResponse brokerResponse = requestHandler.handleRequest( "SELECT now() AS currentTs, fromDateTime('2020-01-01 UTC', 'yyyy-MM-dd z') AS firstDayOf2020"); @@ -349,7 +352,8 @@ public void testExplainPlanLiteralOnly() throws Exception { SingleConnectionBrokerRequestHandler requestHandler = new SingleConnectionBrokerRequestHandler(new PinotConfiguration(), "testBrokerId", null, ACCESS_CONTROL_FACTORY, - null, null, null, null, mock(ServerRoutingStatsManager.class), mock(FailureDetector.class)); + null, null, null, null, mock(ServerRoutingStatsManager.class), mock(FailureDetector.class), + new Tracing.DefaultThreadResourceUsageAccountant()); // Test 1: select constant BrokerResponse brokerResponse = requestHandler.handleRequest("EXPLAIN PLAN FOR SELECT 1.5, 'test'"); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/BrokerReduceService.java b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/BrokerReduceService.java index 9cbd664d53e1..781dbd9b1dfd 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/BrokerReduceService.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/BrokerReduceService.java @@ -36,10 +36,12 @@ import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils; import org.apache.pinot.core.transport.ServerRoutingInstance; import org.apache.pinot.core.util.GapfillUtils; +import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant; import org.apache.pinot.spi.env.PinotConfiguration; import org.apache.pinot.spi.exception.BadQueryRequestException; import org.apache.pinot.spi.exception.EarlyTerminationException; import org.apache.pinot.spi.exception.QueryErrorCode; +import org.apache.pinot.spi.trace.Tracing; import org.apache.pinot.spi.utils.CommonConstants; import org.apache.pinot.spi.utils.builder.TableNameBuilder; import org.slf4j.Logger; @@ -54,8 +56,15 @@ public class BrokerReduceService extends BaseReduceService { private static final Logger LOGGER = LoggerFactory.getLogger(BrokerReduceService.class); + private final ThreadResourceUsageAccountant _resourceUsageAccountant; + public BrokerReduceService(PinotConfiguration config) { + this(config, new Tracing.DefaultThreadResourceUsageAccountant()); + } + + public BrokerReduceService(PinotConfiguration config, ThreadResourceUsageAccountant resourceUsageAccountant) { super(config); + _resourceUsageAccountant = resourceUsageAccountant; } public BrokerResponseNative reduceOnDataTable(BrokerRequest brokerRequest, BrokerRequest serverBrokerRequest, @@ -139,7 +148,8 @@ public BrokerResponseNative reduceOnDataTable(BrokerRequest brokerRequest, Broke } QueryContext serverQueryContext = QueryContextConverterUtils.getQueryContext(serverBrokerRequest.getPinotQuery()); - DataTableReducer dataTableReducer = ResultReducerFactory.getResultReducer(serverQueryContext); + DataTableReducer dataTableReducer = + ResultReducerFactory.getResultReducer(serverQueryContext, _resourceUsageAccountant); Integer minGroupTrimSizeQueryOption = null; Integer groupTrimThresholdQueryOption = null; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java index c5647c1b2737..d385f29dae18 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java @@ -60,6 +60,7 @@ import org.apache.pinot.core.util.GroupByUtils; import org.apache.pinot.core.util.trace.TraceRunnable; import org.apache.pinot.spi.accounting.ThreadExecutionContext; +import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant; import org.apache.pinot.spi.exception.EarlyTerminationException; import org.apache.pinot.spi.exception.QueryErrorCode; import org.apache.pinot.spi.trace.Tracing; @@ -79,8 +80,9 @@ public class GroupByDataTableReducer implements DataTableReducer { private final int _numAggregationFunctions; private final int _numGroupByExpressions; private final int _numColumns; + private final ThreadResourceUsageAccountant _resourceUsageAccountant; - public GroupByDataTableReducer(QueryContext queryContext) { + public GroupByDataTableReducer(QueryContext queryContext, ThreadResourceUsageAccountant accountant) { _queryContext = queryContext; _aggregationFunctions = queryContext.getAggregationFunctions(); assert _aggregationFunctions != null; @@ -89,6 +91,7 @@ public GroupByDataTableReducer(QueryContext queryContext) { assert groupByExpressions != null; _numGroupByExpressions = groupByExpressions.size(); _numColumns = _numAggregationFunctions + _numGroupByExpressions; + _resourceUsageAccountant = accountant; } /** @@ -265,7 +268,7 @@ private IndexedTable getIndexedTable(DataSchema dataSchema, Collection // compare result dataTable against nonStreamingResultDataTable // Process server response. QueryContext queryContext = QueryContextConverterUtils.getQueryContext(sql); - DataTableReducer reducer = ResultReducerFactory.getResultReducer(queryContext); + DataTableReducer reducer = + ResultReducerFactory.getResultReducer(queryContext, new Tracing.DefaultThreadResourceUsageAccountant()); BrokerResponseNative streamingBrokerResponse = new BrokerResponseNative(); reducer.reduceAndSetResults("mytable_OFFLINE", cachedDataSchema, dataTableMap, streamingBrokerResponse, DATATABLE_REDUCER_CONTEXT, mock(BrokerMetrics.class)); diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/server/QueryServer.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/server/QueryServer.java index d6e10f6e60b2..7fa9f157b961 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/server/QueryServer.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/server/QueryServer.java @@ -284,11 +284,11 @@ private void submitInternal(Worker.QueryRequest request, Map req /// (normally cancelling other already started workers and sending the error through GRPC) private CompletableFuture submitWorker(WorkerMetadata workerMetadata, StagePlan stagePlan, Map reqMetadata) { - String requestIdStr = Long.toString(QueryThreadContext.getRequestId()); String workloadName = reqMetadata.get(CommonConstants.Broker.Request.QueryOptionKey.WORKLOAD_NAME); //TODO: Verify if this matches with what OOM protection expects. This method will not block for the query to // finish, so it may be breaking some of the OOM protection assumptions. - Tracing.ThreadAccountantOps.setupRunner(requestIdStr, ThreadExecutionContext.TaskType.MSE, workloadName); + Tracing.ThreadAccountantOps.setupRunner(QueryThreadContext.getCid(), ThreadExecutionContext.TaskType.MSE, + workloadName); ThreadExecutionContext parentContext = Tracing.getThreadAccountant().getThreadExecutionContext(); try { From 92a16d4ecbcf0e395a01aa2b1ab9da0ed6c2314f Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 24 Jul 2025 22:09:49 +0530 Subject: [PATCH 018/167] Bump software.amazon.awssdk:bom from 2.32.6 to 2.32.7 (#16417) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index d2024d480647..c002678b5de6 100644 --- a/pom.xml +++ b/pom.xml @@ -178,7 +178,7 @@ 0.15.1 0.4.7 4.2.2 - 2.32.6 + 2.32.7 1.2.36 1.22.0 2.14.0 From e8b1c52641fe373e2718dbeeca9ae60dabbfaff6 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 24 Jul 2025 22:09:57 +0530 Subject: [PATCH 019/167] Bump curator.version from 5.8.0 to 5.9.0 (#16418) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index c002678b5de6..ff94792856c9 100644 --- a/pom.xml +++ b/pom.xml @@ -261,7 +261,7 @@ 3.6.3 9.4.57.v20241219 7.1.1 - 5.8.0 + 5.9.0 3.30.2-GA 1.78.1 0.27 From 4935e417f3a8bcb5731d428d82fb8aceda6132ec Mon Sep 17 00:00:00 2001 From: Shaurya Chaturvedi Date: Thu, 24 Jul 2025 09:53:50 -0700 Subject: [PATCH 020/167] [timeseries] Adding error banner for timeseries error handling in Controller UI (#16406) Co-authored-by: Shaurya Chaturvedi --- .../components/Query/TimeseriesQueryPage.tsx | 21 ++++++++++---- .../pinot/tsdb/m3ql/parser/Tokenizer.java | 29 +++++++++++++++++++ 2 files changed, 44 insertions(+), 6 deletions(-) diff --git a/pinot-controller/src/main/resources/app/components/Query/TimeseriesQueryPage.tsx b/pinot-controller/src/main/resources/app/components/Query/TimeseriesQueryPage.tsx index 03922a0bc2eb..964945813beb 100644 --- a/pinot-controller/src/main/resources/app/components/Query/TimeseriesQueryPage.tsx +++ b/pinot-controller/src/main/resources/app/components/Query/TimeseriesQueryPage.tsx @@ -51,6 +51,7 @@ import { MAX_SERIES_LIMIT } from '../../utils/ChartConstants'; // Define proper types interface TimeseriesQueryResponse { + error: string; data: { resultType: string; result: Array<{ @@ -335,6 +336,14 @@ const TimeseriesQueryPage = () => { setRawData(parsedData); setRawOutput(JSON.stringify(parsedData, null, 2)); + // Check if this is an error response + if (parsedData.error != null && parsedData.error !== '') { + setError(parsedData.error); + setChartSeries([]); + setTruncatedChartSeries([]); + return; + } + // Parse timeseries data for chart and stats if (isPrometheusFormat(parsedData)) { const series = parseTimeseriesResponse(parsedData); @@ -498,12 +507,6 @@ const TimeseriesQueryPage = () => { - {error && ( - - {error} - - )} - {rawOutput && ( { classes={classes} /> + {error && ( + + {error} + + )} + {viewType === 'chart' && ( > tokenize() { String[] pipelines = _query.split("\\|"); List> result = new ArrayList<>(); for (String pipeline : pipelines) { + Preconditions.checkState(isValidToken(pipeline), String.format("Invalid token: %s", pipeline)); String command = pipeline.trim().substring(0, pipeline.indexOf("{")); if (command.equals("fetch")) { result.add(consumeFetch(pipeline.trim())); @@ -87,4 +88,32 @@ private List consumeGeneric(String pipeline) { } return result; } + + public static boolean isValidToken(String input) { + if (input == null || input.length() < 2) { + return false; + } + int openIndex = -1; + int closeIndex = -1; + // Find first '{' from the front + for (int i = 0; i < input.length(); i++) { + if (input.charAt(i) == '{') { + openIndex = i; + break; + } + } + // If no '{' found + if (openIndex == -1) { + return false; + } + // Find first '}' from the back + for (int i = input.length() - 1; i > openIndex; i--) { + if (input.charAt(i) == '}') { + closeIndex = i; + break; + } + } + // Valid only if '}' found after '{' + return closeIndex != -1; + } } From bd333b4ff73e0370eb23ebad8923601ae849b75e Mon Sep 17 00:00:00 2001 From: "Xiaotian (Jackie) Jiang" <17555551+Jackie-Jiang@users.noreply.github.com> Date: Thu, 24 Jul 2025 16:30:51 -0600 Subject: [PATCH 021/167] Cleanup stream message related APIs (#16388) --- .../realtime/IngestionDelayTracker.java | 8 +- .../realtime/RealtimeSegmentDataManager.java | 21 ++-- .../realtime/RealtimeTableDataManager.java | 8 +- .../kafka20/KafkaStreamMetadataProvider.java | 4 +- .../KafkaPartitionLevelConsumerTest.java | 5 - .../kafka30/KafkaStreamMetadataProvider.java | 4 +- .../KafkaPartitionLevelConsumerTest.java | 5 - .../kafka/KafkaStreamMessageMetadata.java | 15 +-- .../stream/kinesis/KinesisConsumer.java | 4 +- .../kinesis/KinesisStreamMessageMetadata.java | 21 +--- .../KinesisStreamMetadataProvider.java | 4 +- .../kinesis/KinesisMessageBatchTest.java | 1 - .../pulsar/PulsarPartitionLevelConsumer.java | 6 +- .../pulsar/PulsarStreamMessageMetadata.java | 21 +--- .../stream/pulsar/PulsarConsumerTest.java | 4 - .../mutable/MutableSegmentImpl.java | 8 +- .../StatelessRealtimeSegmentWriter.java | 7 +- .../mutable/IndexingFailureTest.java | 14 +-- ...tableSegmentEntriesAboveThresholdTest.java | 12 +-- ...utableSegmentImplAggregateMetricsTest.java | 6 +- ...leSegmentImplIngestionAggregationTest.java | 19 ++-- .../mutable/MutableSegmentImplRawMVTest.java | 13 ++- .../mutable/MutableSegmentImplTest.java | 17 +-- .../pinot/segment/spi/MutableSegment.java | 6 +- .../pinot/spi/stream/BytesStreamMessage.java | 8 +- .../spi/stream/ConsumerPartitionState.java | 7 +- .../apache/pinot/spi/stream/MessageBatch.java | 49 +-------- .../apache/pinot/spi/stream/RowMetadata.java | 102 ------------------ .../spi/stream/StreamDataDecoderImpl.java | 18 ++-- .../pinot/spi/stream/StreamMessage.java | 23 +--- .../spi/stream/StreamMessageMetadata.java | 51 ++------- .../spi/stream/StreamDataDecoderImplTest.java | 29 ++--- 32 files changed, 135 insertions(+), 385 deletions(-) delete mode 100644 pinot-spi/src/main/java/org/apache/pinot/spi/stream/RowMetadata.java diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/IngestionDelayTracker.java b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/IngestionDelayTracker.java index a4cbef048888..90e19ead18f7 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/IngestionDelayTracker.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/IngestionDelayTracker.java @@ -39,7 +39,7 @@ import org.apache.pinot.common.metrics.ServerMetrics; import org.apache.pinot.common.utils.LLCSegmentName; import org.apache.pinot.spi.stream.LongMsgOffset; -import org.apache.pinot.spi.stream.RowMetadata; +import org.apache.pinot.spi.stream.StreamMessageMetadata; import org.apache.pinot.spi.stream.StreamMetadataProvider; import org.apache.pinot.spi.stream.StreamPartitionMsgOffset; import org.apache.pinot.spi.utils.builder.TableNameBuilder; @@ -253,10 +253,10 @@ void setClock(Clock clock) { * * @param segmentName name of the consuming segment * @param partitionId partition id of the consuming segment (directly passed in to avoid parsing the segment name) - * @param ingestionTimeMs ingestion time of the last consumed message (from {@link RowMetadata}) + * @param ingestionTimeMs ingestion time of the last consumed message (from {@link StreamMessageMetadata}) * @param firstStreamIngestionTimeMs ingestion time of the last consumed message in the first stream (from - * {@link RowMetadata}) - * @param currentOffset offset of the last consumed message (from {@link RowMetadata}) + * {@link StreamMessageMetadata}) + * @param currentOffset offset of the last consumed message (from {@link StreamMessageMetadata}) * @param latestOffset offset of the latest message in the partition (from {@link StreamMetadataProvider}) */ public void updateIngestionMetrics(String segmentName, int partitionId, long ingestionTimeMs, diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManager.java b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManager.java index c9ea50f7a845..9c08c0ea4aba 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManager.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManager.java @@ -95,7 +95,6 @@ import org.apache.pinot.spi.stream.PartitionGroupConsumptionStatus; import org.apache.pinot.spi.stream.PartitionLagState; import org.apache.pinot.spi.stream.PermanentConsumerException; -import org.apache.pinot.spi.stream.RowMetadata; import org.apache.pinot.spi.stream.StreamConfig; import org.apache.pinot.spi.stream.StreamConsumerFactory; import org.apache.pinot.spi.stream.StreamConsumerFactoryProvider; @@ -318,7 +317,7 @@ public void deleteSegmentFile() { private final StreamPartitionMsgOffset _startOffset; private final StreamConfig _streamConfig; - private RowMetadata _lastRowMetadata; + private StreamMessageMetadata _lastRowMetadata; private long _lastConsumedTimestampMs = -1; private long _consumeStartTime = -1; private long _lastLogTime = 0; @@ -620,20 +619,12 @@ private boolean processStreamEvents(MessageBatch messageBatch, long idlePipeSlee } // Decode message - StreamMessage streamMessage = messageBatch.getStreamMessage(index); + StreamMessage streamMessage = messageBatch.getStreamMessage(index); StreamDataDecoderResult decodedRow = _streamDataDecoder.decode(streamMessage); StreamMessageMetadata metadata = streamMessage.getMetadata(); - StreamPartitionMsgOffset offset = null; - StreamPartitionMsgOffset nextOffset = null; - if (metadata != null) { - offset = metadata.getOffset(); - nextOffset = metadata.getNextOffset(); - } - // Backward compatible - if (nextOffset == null) { - nextOffset = messageBatch.getNextStreamPartitionMsgOffsetAtIndex(index); - } - int rowSizeInBytes = null == metadata ? 0 : metadata.getRecordSerializedSize(); + StreamPartitionMsgOffset offset = metadata.getOffset(); + StreamPartitionMsgOffset nextOffset = metadata.getNextOffset(); + int rowSizeInBytes = metadata.getRecordSerializedSize(); if (decodedRow.getException() != null) { // TODO: based on a config, decide whether the record should be silently dropped or stop further consumption on // decode error @@ -2018,7 +2009,7 @@ private void createPartitionMetadataProvider(String reason) { _streamConsumerFactory.createPartitionMetadataProvider(_clientId, _streamPartitionId); } - private void updateIngestionMetrics(RowMetadata metadata) { + private void updateIngestionMetrics(StreamMessageMetadata metadata) { if (metadata != null) { try { StreamPartitionMsgOffset latestOffset = fetchLatestStreamOffset(5000, true); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeTableDataManager.java b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeTableDataManager.java index e25f8da05e2c..ed80b37abe6d 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeTableDataManager.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeTableDataManager.java @@ -82,8 +82,8 @@ import org.apache.pinot.spi.data.DateTimeFormatSpec; import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.data.Schema; -import org.apache.pinot.spi.stream.RowMetadata; import org.apache.pinot.spi.stream.StreamConfigProperties; +import org.apache.pinot.spi.stream.StreamMessageMetadata; import org.apache.pinot.spi.stream.StreamMetadataProvider; import org.apache.pinot.spi.stream.StreamPartitionMsgOffset; import org.apache.pinot.spi.utils.CommonConstants; @@ -293,10 +293,10 @@ protected void doShutdown() { * * @param segmentName name of the consuming segment * @param partitionId partition id of the consuming segment (directly passed in to avoid parsing the segment name) - * @param ingestionTimeMs ingestion time of the last consumed message (from {@link RowMetadata}) + * @param ingestionTimeMs ingestion time of the last consumed message (from {@link StreamMessageMetadata}) * @param firstStreamIngestionTimeMs ingestion time of the last consumed message in the first stream (from - * {@link RowMetadata}) - * @param currentOffset offset of the last consumed message (from {@link RowMetadata}) + * {@link StreamMessageMetadata}) + * @param currentOffset offset of the last consumed message (from {@link StreamMessageMetadata}) * @param latestOffset offset of the latest message in the partition (from {@link StreamMetadataProvider}) */ public void updateIngestionMetrics(String segmentName, int partitionId, long ingestionTimeMs, diff --git a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-2.0/src/main/java/org/apache/pinot/plugin/stream/kafka20/KafkaStreamMetadataProvider.java b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-2.0/src/main/java/org/apache/pinot/plugin/stream/kafka20/KafkaStreamMetadataProvider.java index a04cca66d2a1..0b13c05ba482 100644 --- a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-2.0/src/main/java/org/apache/pinot/plugin/stream/kafka20/KafkaStreamMetadataProvider.java +++ b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-2.0/src/main/java/org/apache/pinot/plugin/stream/kafka20/KafkaStreamMetadataProvider.java @@ -41,8 +41,8 @@ import org.apache.pinot.spi.stream.LongMsgOffset; import org.apache.pinot.spi.stream.OffsetCriteria; import org.apache.pinot.spi.stream.PartitionLagState; -import org.apache.pinot.spi.stream.RowMetadata; import org.apache.pinot.spi.stream.StreamConfig; +import org.apache.pinot.spi.stream.StreamMessageMetadata; import org.apache.pinot.spi.stream.StreamMetadataProvider; import org.apache.pinot.spi.stream.StreamPartitionMsgOffset; import org.apache.pinot.spi.stream.TransientConsumerException; @@ -158,7 +158,7 @@ public Map getCurrentPartitionLagState( // Compute record-availability String availabilityLagMs = "UNKNOWN"; - RowMetadata lastProcessedMessageMetadata = partitionState.getLastProcessedRowMetadata(); + StreamMessageMetadata lastProcessedMessageMetadata = partitionState.getLastProcessedRowMetadata(); if (lastProcessedMessageMetadata != null && partitionState.getLastProcessedTimeMs() > 0) { long availabilityLag = partitionState.getLastProcessedTimeMs() - lastProcessedMessageMetadata.getRecordIngestionTimeMs(); diff --git a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-2.0/src/test/java/org/apache/pinot/plugin/stream/kafka20/KafkaPartitionLevelConsumerTest.java b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-2.0/src/test/java/org/apache/pinot/plugin/stream/kafka20/KafkaPartitionLevelConsumerTest.java index 4cf1be35dfe3..5a25ef195053 100644 --- a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-2.0/src/test/java/org/apache/pinot/plugin/stream/kafka20/KafkaPartitionLevelConsumerTest.java +++ b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-2.0/src/test/java/org/apache/pinot/plugin/stream/kafka20/KafkaPartitionLevelConsumerTest.java @@ -48,7 +48,6 @@ import org.testng.annotations.Test; import static org.testng.Assert.assertEquals; -import static org.testng.Assert.assertNotNull; import static org.testng.Assert.assertTrue; @@ -278,7 +277,6 @@ private void testConsumer(String topic) StreamMessage streamMessage = messageBatch.getStreamMessage(i); assertEquals(new String((byte[]) streamMessage.getValue()), "sample_msg_" + i); StreamMessageMetadata metadata = streamMessage.getMetadata(); - assertNotNull(metadata); assertEquals(metadata.getRecordIngestionTimeMs(), TIMESTAMP + i); StreamPartitionMsgOffset offset = metadata.getOffset(); assertTrue(offset instanceof LongMsgOffset); @@ -300,7 +298,6 @@ private void testConsumer(String topic) StreamMessage streamMessage = messageBatch.getStreamMessage(i); assertEquals(new String((byte[]) streamMessage.getValue()), "sample_msg_" + (500 + i)); StreamMessageMetadata metadata = streamMessage.getMetadata(); - assertNotNull(metadata); assertEquals(metadata.getRecordIngestionTimeMs(), TIMESTAMP + 500 + i); StreamPartitionMsgOffset offset = metadata.getOffset(); assertTrue(offset instanceof LongMsgOffset); @@ -322,7 +319,6 @@ private void testConsumer(String topic) StreamMessage streamMessage = messageBatch.getStreamMessage(i); assertEquals(new String((byte[]) streamMessage.getValue()), "sample_msg_" + (10 + i)); StreamMessageMetadata metadata = streamMessage.getMetadata(); - assertNotNull(metadata); assertEquals(metadata.getRecordIngestionTimeMs(), TIMESTAMP + 10 + i); StreamPartitionMsgOffset offset = metadata.getOffset(); assertTrue(offset instanceof LongMsgOffset); @@ -344,7 +340,6 @@ private void testConsumer(String topic) StreamMessage streamMessage = messageBatch.getStreamMessage(i); assertEquals(new String((byte[]) streamMessage.getValue()), "sample_msg_" + (610 + i)); StreamMessageMetadata metadata = streamMessage.getMetadata(); - assertNotNull(metadata); assertEquals(metadata.getRecordIngestionTimeMs(), TIMESTAMP + 610 + i); StreamPartitionMsgOffset offset = metadata.getOffset(); assertTrue(offset instanceof LongMsgOffset); diff --git a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/main/java/org/apache/pinot/plugin/stream/kafka30/KafkaStreamMetadataProvider.java b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/main/java/org/apache/pinot/plugin/stream/kafka30/KafkaStreamMetadataProvider.java index 96775641ca31..62a181f605dd 100644 --- a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/main/java/org/apache/pinot/plugin/stream/kafka30/KafkaStreamMetadataProvider.java +++ b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/main/java/org/apache/pinot/plugin/stream/kafka30/KafkaStreamMetadataProvider.java @@ -41,8 +41,8 @@ import org.apache.pinot.spi.stream.LongMsgOffset; import org.apache.pinot.spi.stream.OffsetCriteria; import org.apache.pinot.spi.stream.PartitionLagState; -import org.apache.pinot.spi.stream.RowMetadata; import org.apache.pinot.spi.stream.StreamConfig; +import org.apache.pinot.spi.stream.StreamMessageMetadata; import org.apache.pinot.spi.stream.StreamMetadataProvider; import org.apache.pinot.spi.stream.StreamPartitionMsgOffset; import org.apache.pinot.spi.stream.TransientConsumerException; @@ -158,7 +158,7 @@ public Map getCurrentPartitionLagState( // Compute record-availability String availabilityLagMs = "UNKNOWN"; - RowMetadata lastProcessedMessageMetadata = partitionState.getLastProcessedRowMetadata(); + StreamMessageMetadata lastProcessedMessageMetadata = partitionState.getLastProcessedRowMetadata(); if (lastProcessedMessageMetadata != null && partitionState.getLastProcessedTimeMs() > 0) { long availabilityLag = partitionState.getLastProcessedTimeMs() - lastProcessedMessageMetadata.getRecordIngestionTimeMs(); diff --git a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/test/java/org/apache/pinot/plugin/stream/kafka30/KafkaPartitionLevelConsumerTest.java b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/test/java/org/apache/pinot/plugin/stream/kafka30/KafkaPartitionLevelConsumerTest.java index 1c52b37f3e55..6df03f8dbceb 100644 --- a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/test/java/org/apache/pinot/plugin/stream/kafka30/KafkaPartitionLevelConsumerTest.java +++ b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-3.0/src/test/java/org/apache/pinot/plugin/stream/kafka30/KafkaPartitionLevelConsumerTest.java @@ -48,7 +48,6 @@ import org.testng.annotations.Test; import static org.testng.Assert.assertEquals; -import static org.testng.Assert.assertNotNull; import static org.testng.Assert.assertTrue; @@ -278,7 +277,6 @@ private void testConsumer(String topic) StreamMessage streamMessage = messageBatch.getStreamMessage(i); assertEquals(new String((byte[]) streamMessage.getValue()), "sample_msg_" + i); StreamMessageMetadata metadata = streamMessage.getMetadata(); - assertNotNull(metadata); assertEquals(metadata.getRecordIngestionTimeMs(), TIMESTAMP + i); StreamPartitionMsgOffset offset = metadata.getOffset(); assertTrue(offset instanceof LongMsgOffset); @@ -300,7 +298,6 @@ private void testConsumer(String topic) StreamMessage streamMessage = messageBatch.getStreamMessage(i); assertEquals(new String((byte[]) streamMessage.getValue()), "sample_msg_" + (500 + i)); StreamMessageMetadata metadata = streamMessage.getMetadata(); - assertNotNull(metadata); assertEquals(metadata.getRecordIngestionTimeMs(), TIMESTAMP + 500 + i); StreamPartitionMsgOffset offset = metadata.getOffset(); assertTrue(offset instanceof LongMsgOffset); @@ -322,7 +319,6 @@ private void testConsumer(String topic) StreamMessage streamMessage = messageBatch.getStreamMessage(i); assertEquals(new String((byte[]) streamMessage.getValue()), "sample_msg_" + (10 + i)); StreamMessageMetadata metadata = streamMessage.getMetadata(); - assertNotNull(metadata); assertEquals(metadata.getRecordIngestionTimeMs(), TIMESTAMP + 10 + i); StreamPartitionMsgOffset offset = metadata.getOffset(); assertTrue(offset instanceof LongMsgOffset); @@ -344,7 +340,6 @@ private void testConsumer(String topic) StreamMessage streamMessage = messageBatch.getStreamMessage(i); assertEquals(new String((byte[]) streamMessage.getValue()), "sample_msg_" + (610 + i)); StreamMessageMetadata metadata = streamMessage.getMetadata(); - assertNotNull(metadata); assertEquals(metadata.getRecordIngestionTimeMs(), TIMESTAMP + 610 + i); StreamPartitionMsgOffset offset = metadata.getOffset(); assertTrue(offset instanceof LongMsgOffset); diff --git a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-base/src/main/java/org/apache/pinot/plugin/stream/kafka/KafkaStreamMessageMetadata.java b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-base/src/main/java/org/apache/pinot/plugin/stream/kafka/KafkaStreamMessageMetadata.java index f3087a879ca0..17d1fbc1d2d0 100644 --- a/pinot-plugins/pinot-stream-ingestion/pinot-kafka-base/src/main/java/org/apache/pinot/plugin/stream/kafka/KafkaStreamMessageMetadata.java +++ b/pinot-plugins/pinot-stream-ingestion/pinot-kafka-base/src/main/java/org/apache/pinot/plugin/stream/kafka/KafkaStreamMessageMetadata.java @@ -18,20 +18,11 @@ */ package org.apache.pinot.plugin.stream.kafka; -import java.util.Map; -import javax.annotation.Nullable; -import org.apache.pinot.spi.data.readers.GenericRow; -import org.apache.pinot.spi.stream.StreamMessageMetadata; +public class KafkaStreamMessageMetadata { + private KafkaStreamMessageMetadata() { + } -// TODO: Make it a util class -public class KafkaStreamMessageMetadata extends StreamMessageMetadata { public static final String METADATA_OFFSET_KEY = "offset"; public static final String RECORD_TIMESTAMP_KEY = "recordTimestamp"; public static final String METADATA_PARTITION_KEY = "partition"; - - @Deprecated - public KafkaStreamMessageMetadata(long recordIngestionTimeMs, @Nullable GenericRow headers, - Map metadata) { - super(recordIngestionTimeMs, headers, metadata); - } } diff --git a/pinot-plugins/pinot-stream-ingestion/pinot-kinesis/src/main/java/org/apache/pinot/plugin/stream/kinesis/KinesisConsumer.java b/pinot-plugins/pinot-stream-ingestion/pinot-kinesis/src/main/java/org/apache/pinot/plugin/stream/kinesis/KinesisConsumer.java index f5a905e111ea..5e8f84010762 100644 --- a/pinot-plugins/pinot-stream-ingestion/pinot-kinesis/src/main/java/org/apache/pinot/plugin/stream/kinesis/KinesisConsumer.java +++ b/pinot-plugins/pinot-stream-ingestion/pinot-kinesis/src/main/java/org/apache/pinot/plugin/stream/kinesis/KinesisConsumer.java @@ -114,9 +114,7 @@ private KinesisMessageBatch getKinesisMessageBatch(KinesisPartitionGroupOffset s KinesisPartitionGroupOffset offsetOfNextBatch; if (!records.isEmpty()) { messages = records.stream().map(record -> extractStreamMessage(record, shardId)).collect(Collectors.toList()); - StreamMessageMetadata lastMessageMetadata = messages.get(messages.size() - 1).getMetadata(); - assert lastMessageMetadata != null; - offsetOfNextBatch = (KinesisPartitionGroupOffset) lastMessageMetadata.getNextOffset(); + offsetOfNextBatch = (KinesisPartitionGroupOffset) messages.get(messages.size() - 1).getMetadata().getNextOffset(); } else { // TODO: Revisit whether Kinesis can return empty batch when there are available records. The consumer cna handle // empty message batch, but it will treat it as fully caught up. diff --git a/pinot-plugins/pinot-stream-ingestion/pinot-kinesis/src/main/java/org/apache/pinot/plugin/stream/kinesis/KinesisStreamMessageMetadata.java b/pinot-plugins/pinot-stream-ingestion/pinot-kinesis/src/main/java/org/apache/pinot/plugin/stream/kinesis/KinesisStreamMessageMetadata.java index cbabaf608a09..0ba82337568c 100644 --- a/pinot-plugins/pinot-stream-ingestion/pinot-kinesis/src/main/java/org/apache/pinot/plugin/stream/kinesis/KinesisStreamMessageMetadata.java +++ b/pinot-plugins/pinot-stream-ingestion/pinot-kinesis/src/main/java/org/apache/pinot/plugin/stream/kinesis/KinesisStreamMessageMetadata.java @@ -18,25 +18,10 @@ */ package org.apache.pinot.plugin.stream.kinesis; -import java.util.Map; -import javax.annotation.Nullable; -import org.apache.pinot.spi.data.readers.GenericRow; -import org.apache.pinot.spi.stream.StreamMessageMetadata; - +public class KinesisStreamMessageMetadata { + private KinesisStreamMessageMetadata() { + } -// TODO: Make it a util class -public class KinesisStreamMessageMetadata extends StreamMessageMetadata { public static final String APPRX_ARRIVAL_TIMESTAMP_KEY = "apprxArrivalTimestamp"; public static final String SEQUENCE_NUMBER_KEY = "sequenceNumber"; - - @Deprecated - public KinesisStreamMessageMetadata(long recordIngestionTimeMs, @Nullable GenericRow headers) { - super(recordIngestionTimeMs, headers); - } - - @Deprecated - public KinesisStreamMessageMetadata(long recordIngestionTimeMs, @Nullable GenericRow headers, - Map metadata) { - super(recordIngestionTimeMs, headers, metadata); - } } diff --git a/pinot-plugins/pinot-stream-ingestion/pinot-kinesis/src/main/java/org/apache/pinot/plugin/stream/kinesis/KinesisStreamMetadataProvider.java b/pinot-plugins/pinot-stream-ingestion/pinot-kinesis/src/main/java/org/apache/pinot/plugin/stream/kinesis/KinesisStreamMetadataProvider.java index d9b5f17e3971..fb24821910a7 100644 --- a/pinot-plugins/pinot-stream-ingestion/pinot-kinesis/src/main/java/org/apache/pinot/plugin/stream/kinesis/KinesisStreamMetadataProvider.java +++ b/pinot-plugins/pinot-stream-ingestion/pinot-kinesis/src/main/java/org/apache/pinot/plugin/stream/kinesis/KinesisStreamMetadataProvider.java @@ -35,10 +35,10 @@ import org.apache.pinot.spi.stream.PartitionGroupConsumptionStatus; import org.apache.pinot.spi.stream.PartitionGroupMetadata; import org.apache.pinot.spi.stream.PartitionLagState; -import org.apache.pinot.spi.stream.RowMetadata; import org.apache.pinot.spi.stream.StreamConfig; import org.apache.pinot.spi.stream.StreamConsumerFactory; import org.apache.pinot.spi.stream.StreamConsumerFactoryProvider; +import org.apache.pinot.spi.stream.StreamMessageMetadata; import org.apache.pinot.spi.stream.StreamMetadataProvider; import org.apache.pinot.spi.stream.StreamPartitionMsgOffset; import org.slf4j.Logger; @@ -264,7 +264,7 @@ public Map getCurrentPartitionLagState( ConsumerPartitionState partitionState = entry.getValue(); // Compute record-availability String recordAvailabilityLag = "UNKNOWN"; - RowMetadata lastProcessedMessageMetadata = partitionState.getLastProcessedRowMetadata(); + StreamMessageMetadata lastProcessedMessageMetadata = partitionState.getLastProcessedRowMetadata(); if (lastProcessedMessageMetadata != null && partitionState.getLastProcessedTimeMs() > 0) { long availabilityLag = partitionState.getLastProcessedTimeMs() - lastProcessedMessageMetadata.getRecordIngestionTimeMs(); diff --git a/pinot-plugins/pinot-stream-ingestion/pinot-kinesis/src/test/java/org/apache/pinot/plugin/stream/kinesis/KinesisMessageBatchTest.java b/pinot-plugins/pinot-stream-ingestion/pinot-kinesis/src/test/java/org/apache/pinot/plugin/stream/kinesis/KinesisMessageBatchTest.java index 4e64dae54961..a727c33b72af 100644 --- a/pinot-plugins/pinot-stream-ingestion/pinot-kinesis/src/test/java/org/apache/pinot/plugin/stream/kinesis/KinesisMessageBatchTest.java +++ b/pinot-plugins/pinot-stream-ingestion/pinot-kinesis/src/test/java/org/apache/pinot/plugin/stream/kinesis/KinesisMessageBatchTest.java @@ -52,7 +52,6 @@ public void testMessageBatch() { byte[] value = streamMessage.getValue(); assertEquals(new String(value, StandardCharsets.UTF_8), "value-" + i); StreamMessageMetadata metadata = streamMessage.getMetadata(); - assertNotNull(metadata); assertEquals(metadata.getRecordIngestionTimeMs(), baseTimeMs + i); StreamPartitionMsgOffset offset = metadata.getOffset(); assertTrue(offset instanceof KinesisPartitionGroupOffset); diff --git a/pinot-plugins/pinot-stream-ingestion/pinot-pulsar/src/main/java/org/apache/pinot/plugin/stream/pulsar/PulsarPartitionLevelConsumer.java b/pinot-plugins/pinot-stream-ingestion/pinot-pulsar/src/main/java/org/apache/pinot/plugin/stream/pulsar/PulsarPartitionLevelConsumer.java index c206574bc924..d96d137eeed4 100644 --- a/pinot-plugins/pinot-stream-ingestion/pinot-pulsar/src/main/java/org/apache/pinot/plugin/stream/pulsar/PulsarPartitionLevelConsumer.java +++ b/pinot-plugins/pinot-stream-ingestion/pinot-pulsar/src/main/java/org/apache/pinot/plugin/stream/pulsar/PulsarPartitionLevelConsumer.java @@ -25,7 +25,6 @@ import org.apache.pinot.spi.stream.BytesStreamMessage; import org.apache.pinot.spi.stream.PartitionGroupConsumer; import org.apache.pinot.spi.stream.StreamConfig; -import org.apache.pinot.spi.stream.StreamMessageMetadata; import org.apache.pinot.spi.stream.StreamPartitionMsgOffset; import org.apache.pulsar.client.api.MessageId; import org.apache.pulsar.client.api.PulsarClientException; @@ -87,11 +86,8 @@ public synchronized PulsarMessageBatch fetchMessages(StreamPartitionMsgOffset st if (messages.isEmpty()) { offsetOfNextBatch = (MessageIdStreamOffset) startOffset; } else { - StreamMessageMetadata lastMessageMetadata = messages.get(messages.size() - 1).getMetadata(); - assert lastMessageMetadata != null; - offsetOfNextBatch = (MessageIdStreamOffset) lastMessageMetadata.getNextOffset(); + offsetOfNextBatch = (MessageIdStreamOffset) messages.get(messages.size() - 1).getMetadata().getNextOffset(); } - assert offsetOfNextBatch != null; _nextMessageId = offsetOfNextBatch.getMessageId(); return new PulsarMessageBatch(messages, offsetOfNextBatch, _reader.hasReachedEndOfTopic()); } diff --git a/pinot-plugins/pinot-stream-ingestion/pinot-pulsar/src/main/java/org/apache/pinot/plugin/stream/pulsar/PulsarStreamMessageMetadata.java b/pinot-plugins/pinot-stream-ingestion/pinot-pulsar/src/main/java/org/apache/pinot/plugin/stream/pulsar/PulsarStreamMessageMetadata.java index fcf219e98df5..6fec7a05dde2 100644 --- a/pinot-plugins/pinot-stream-ingestion/pinot-pulsar/src/main/java/org/apache/pinot/plugin/stream/pulsar/PulsarStreamMessageMetadata.java +++ b/pinot-plugins/pinot-stream-ingestion/pinot-pulsar/src/main/java/org/apache/pinot/plugin/stream/pulsar/PulsarStreamMessageMetadata.java @@ -20,19 +20,15 @@ package org.apache.pinot.plugin.stream.pulsar; import java.util.EnumSet; -import java.util.Map; -import javax.annotation.Nullable; -import org.apache.pinot.spi.data.readers.GenericRow; -import org.apache.pinot.spi.stream.StreamMessageMetadata; /** - * Pulsar specific implementation of {@link StreamMessageMetadata} * Pulsar makes many metadata values available for each message. Please see the pulsar documentation for more details. * @see Pulsar Message Properties */ -// TODO: Make it a util class -public class PulsarStreamMessageMetadata extends StreamMessageMetadata { +public class PulsarStreamMessageMetadata { + private PulsarStreamMessageMetadata() { + } public enum PulsarMessageMetadataValue { PUBLISH_TIME("publishTime"), @@ -65,15 +61,4 @@ public static PulsarMessageMetadataValue findByKey(final String key) { return values.stream().filter(value -> value.getKey().equals(key)).findFirst().orElse(null); } } - - @Deprecated - public PulsarStreamMessageMetadata(long recordIngestionTimeMs, @Nullable GenericRow headers) { - super(recordIngestionTimeMs, headers); - } - - @Deprecated - public PulsarStreamMessageMetadata(long recordIngestionTimeMs, @Nullable GenericRow headers, - Map metadata) { - super(recordIngestionTimeMs, headers, metadata); - } } diff --git a/pinot-plugins/pinot-stream-ingestion/pinot-pulsar/src/test/java/org/apache/pinot/plugin/stream/pulsar/PulsarConsumerTest.java b/pinot-plugins/pinot-stream-ingestion/pinot-pulsar/src/test/java/org/apache/pinot/plugin/stream/pulsar/PulsarConsumerTest.java index a6414fc713b1..7a4c806c7373 100644 --- a/pinot-plugins/pinot-stream-ingestion/pinot-pulsar/src/test/java/org/apache/pinot/plugin/stream/pulsar/PulsarConsumerTest.java +++ b/pinot-plugins/pinot-stream-ingestion/pinot-pulsar/src/test/java/org/apache/pinot/plugin/stream/pulsar/PulsarConsumerTest.java @@ -51,7 +51,6 @@ import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; -import static org.testng.Assert.assertNotNull; import static org.testng.Assert.assertTrue; @@ -244,11 +243,8 @@ private void testConsumer(PulsarPartitionLevelConsumer consumer, int startIndex, private void verifyMessage(BytesStreamMessage streamMessage, int index, List messageIds) { assertEquals(new String(streamMessage.getValue()), MESSAGE_PREFIX + index); StreamMessageMetadata messageMetadata = streamMessage.getMetadata(); - assertNotNull(messageMetadata); MessageIdStreamOffset offset = (MessageIdStreamOffset) messageMetadata.getOffset(); - assertNotNull(offset); MessageIdStreamOffset nextOffset = (MessageIdStreamOffset) messageMetadata.getNextOffset(); - assertNotNull(nextOffset); assertEquals(offset.getMessageId(), messageIds.get(index)); if (index < NUM_RECORDS_PER_PARTITION - 1) { assertEquals(nextOffset.getMessageId(), messageIds.get(index + 1)); diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImpl.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImpl.java index 5a202999ea07..a3a6b31fdd1e 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImpl.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImpl.java @@ -113,7 +113,7 @@ import org.apache.pinot.spi.data.Schema; import org.apache.pinot.spi.data.readers.GenericRow; import org.apache.pinot.spi.data.readers.PrimaryKey; -import org.apache.pinot.spi.stream.RowMetadata; +import org.apache.pinot.spi.stream.StreamMessageMetadata; import org.apache.pinot.spi.utils.BooleanUtils; import org.apache.pinot.spi.utils.ByteArray; import org.apache.pinot.spi.utils.FixedIntArray; @@ -610,7 +610,7 @@ public long getMaxTime() { } @Override - public boolean index(GenericRow row, @Nullable RowMetadata rowMetadata) + public boolean index(GenericRow row, @Nullable StreamMessageMetadata metadata) throws IOException { boolean canTakeMore; int numDocsIndexed = _numDocsIndexed; @@ -689,8 +689,8 @@ public boolean index(GenericRow row, @Nullable RowMetadata rowMetadata) // Update last indexed time and latest ingestion time _lastIndexedTimeMs = System.currentTimeMillis(); - if (rowMetadata != null) { - _latestIngestionTimeMs = Math.max(_latestIngestionTimeMs, rowMetadata.getRecordIngestionTimeMs()); + if (metadata != null) { + _latestIngestionTimeMs = Math.max(_latestIngestionTimeMs, metadata.getRecordIngestionTimeMs()); } return canTakeMore; diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/writer/StatelessRealtimeSegmentWriter.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/writer/StatelessRealtimeSegmentWriter.java index 05857b113ab9..c56d97e64032 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/writer/StatelessRealtimeSegmentWriter.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/writer/StatelessRealtimeSegmentWriter.java @@ -63,6 +63,7 @@ import org.apache.pinot.spi.stream.StreamDataDecoderResult; import org.apache.pinot.spi.stream.StreamMessage; import org.apache.pinot.spi.stream.StreamMessageDecoder; +import org.apache.pinot.spi.stream.StreamMessageMetadata; import org.apache.pinot.spi.stream.StreamMetadataProvider; import org.apache.pinot.spi.stream.StreamPartitionMsgOffset; import org.apache.pinot.spi.stream.StreamPartitionMsgOffsetFactory; @@ -263,8 +264,8 @@ public void run() { for (int i = 0; i < messageCount; i++) { StreamMessage streamMessage = messageBatch.getStreamMessage(i); - if (streamMessage.getMetadata() != null && streamMessage.getMetadata().getOffset() != null - && streamMessage.getMetadata().getOffset().compareTo(_endOffset) >= 0) { + StreamMessageMetadata metadata = streamMessage.getMetadata(); + if (metadata.getOffset().compareTo(_endOffset) >= 0) { _logger.info("Reached end offset: {} for partition group: {}", _endOffset, _partitionGroupId); break; } @@ -277,7 +278,7 @@ public void run() { assert row != null; TransformPipeline.Result result = _transformPipeline.processRow(row); for (GenericRow transformedRow : result.getTransformedRows()) { - _realtimeSegment.index(transformedRow, streamMessage.getMetadata()); + _realtimeSegment.index(transformedRow, metadata); } } else { _logger.warn("Failed to decode message at offset {}: {}", _currentOffset, decodedResult.getException()); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/IndexingFailureTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/IndexingFailureTest.java index 86dfcbf27643..5e905a371da2 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/IndexingFailureTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/IndexingFailureTest.java @@ -47,6 +47,7 @@ public class IndexingFailureTest implements PinotBuffersAfterMethodCheckRule { private static final String INT_COL = "int_col"; private static final String STRING_COL = "string_col"; private static final String JSON_COL = "json_col"; + private static final StreamMessageMetadata METADATA = mock(StreamMessageMetadata.class); private MutableSegmentImpl _mutableSegment; private ServerMetrics _serverMetrics; @@ -55,7 +56,9 @@ public class IndexingFailureTest implements PinotBuffersAfterMethodCheckRule { public void setup() { Schema schema = new Schema.SchemaBuilder().addSingleValueDimension(INT_COL, FieldSpec.DataType.INT) .addSingleValueDimension(STRING_COL, FieldSpec.DataType.STRING) - .addSingleValueDimension(JSON_COL, FieldSpec.DataType.JSON).setSchemaName(TABLE_NAME).build(); + .addSingleValueDimension(JSON_COL, FieldSpec.DataType.JSON) + .setSchemaName(TABLE_NAME) + .build(); _serverMetrics = mock(ServerMetrics.class); _mutableSegment = MutableSegmentImplTestUtils.createMutableSegmentImpl(schema, Collections.emptySet(), Collections.emptySet(), @@ -71,12 +74,11 @@ public void tearDown() { @Test public void testIndexingFailures() throws IOException { - StreamMessageMetadata defaultMetadata = new StreamMessageMetadata(System.currentTimeMillis(), new GenericRow()); GenericRow goodRow = new GenericRow(); goodRow.putValue(INT_COL, 0); goodRow.putValue(STRING_COL, "a"); goodRow.putValue(JSON_COL, "{\"valid\": \"json\"}"); - _mutableSegment.index(goodRow, defaultMetadata); + _mutableSegment.index(goodRow, METADATA); assertEquals(_mutableSegment.getNumDocsIndexed(), 1); assertEquals(_mutableSegment.getDataSource(INT_COL).getInvertedIndex().getDocIds(0), ImmutableRoaringBitmap.bitmapOf(0)); @@ -92,7 +94,7 @@ public void testIndexingFailures() badRow.putValue(INT_COL, 0); badRow.putValue(STRING_COL, "b"); badRow.putValue(JSON_COL, "{\"truncatedJson..."); - _mutableSegment.index(badRow, defaultMetadata); + _mutableSegment.index(badRow, METADATA); assertEquals(_mutableSegment.getNumDocsIndexed(), 2); assertEquals(_mutableSegment.getDataSource(INT_COL).getInvertedIndex().getDocIds(0), ImmutableRoaringBitmap.bitmapOf(0, 1)); @@ -106,7 +108,7 @@ public void testIndexingFailures() anotherGoodRow.putValue(INT_COL, 2); anotherGoodRow.putValue(STRING_COL, "c"); anotherGoodRow.putValue(JSON_COL, "{\"valid\": \"json\"}"); - _mutableSegment.index(anotherGoodRow, defaultMetadata); + _mutableSegment.index(anotherGoodRow, METADATA); assertEquals(_mutableSegment.getNumDocsIndexed(), 3); assertEquals(_mutableSegment.getDataSource(INT_COL).getInvertedIndex().getDocIds(1), ImmutableRoaringBitmap.bitmapOf(2)); @@ -123,7 +125,7 @@ public void testIndexingFailures() nullStringRow.putValue(STRING_COL, null); nullStringRow.addNullValueField(STRING_COL); nullStringRow.putValue(JSON_COL, "{\"valid\": \"json\"}"); - _mutableSegment.index(nullStringRow, defaultMetadata); + _mutableSegment.index(nullStringRow, METADATA); assertEquals(_mutableSegment.getNumDocsIndexed(), 4); assertEquals(_mutableSegment.getDataSource(INT_COL).getInvertedIndex().getDocIds(0), ImmutableRoaringBitmap.bitmapOf(0, 1, 3)); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentEntriesAboveThresholdTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentEntriesAboveThresholdTest.java index 9eea0635459f..927f41ea3b1f 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentEntriesAboveThresholdTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentEntriesAboveThresholdTest.java @@ -49,6 +49,7 @@ import org.testng.Assert; import org.testng.annotations.Test; +import static org.mockito.Mockito.mock; import static org.mockito.Mockito.spy; @@ -56,6 +57,7 @@ public class MutableSegmentEntriesAboveThresholdTest implements PinotBuffersAfte private static final File TEMP_DIR = new File(FileUtils.getTempDirectory(), MutableSegmentEntriesAboveThresholdTest.class.getSimpleName()); private static final String AVRO_FILE = "data/test_data-mv.avro"; + private static final StreamMessageMetadata METADATA = mock(StreamMessageMetadata.class); private Schema _schema; private File getAvroFile() { @@ -85,12 +87,11 @@ public void testNoLimitBreached() File avroFile = getAvroFile(); MutableSegmentImpl mutableSegment = getMutableSegment(avroFile); try { - StreamMessageMetadata defaultMetadata = new StreamMessageMetadata(System.currentTimeMillis(), new GenericRow()); try (RecordReader recordReader = RecordReaderFactory .getRecordReader(FileFormat.AVRO, avroFile, _schema.getColumnNames(), null)) { GenericRow reuse = new GenericRow(); while (recordReader.hasNext()) { - mutableSegment.index(recordReader.next(reuse), defaultMetadata); + mutableSegment.index(recordReader.next(reuse), METADATA); } } assert mutableSegment.canAddMore(); @@ -105,7 +106,6 @@ public void testLimitBreachedByMutableForwardIndex() File avroFile = getAvroFile(); MutableSegmentImpl mutableSegment = getMutableSegment(avroFile); try { - Field indexContainerMapField = MutableSegmentImpl.class.getDeclaredField("_indexContainerMap"); indexContainerMapField.setAccessible(true); Map colVsIndexContainer = (Map) indexContainerMapField.get(mutableSegment); @@ -129,12 +129,11 @@ public void testLimitBreachedByMutableForwardIndex() indexTypeVsMutableIndex.put(new ForwardIndexPlugin().getIndexType(), new FakeMutableForwardIndex(mutableForwardIndex)); } - StreamMessageMetadata defaultMetadata = new StreamMessageMetadata(System.currentTimeMillis(), new GenericRow()); try (RecordReader recordReader = RecordReaderFactory .getRecordReader(FileFormat.AVRO, avroFile, _schema.getColumnNames(), null)) { GenericRow reuse = new GenericRow(); while (recordReader.hasNext()) { - mutableSegment.index(recordReader.next(reuse), defaultMetadata); + mutableSegment.index(recordReader.next(reuse), METADATA); } } @@ -166,12 +165,11 @@ public void testLimitBreachedByMutableDictionary() mutableDictinaryField.set(indexContainer, mockedDictionary); } } - StreamMessageMetadata defaultMetadata = new StreamMessageMetadata(System.currentTimeMillis(), new GenericRow()); try (RecordReader recordReader = RecordReaderFactory .getRecordReader(FileFormat.AVRO, avroFile, _schema.getColumnNames(), null)) { GenericRow reuse = new GenericRow(); while (recordReader.hasNext()) { - mutableSegment.index(recordReader.next(reuse), defaultMetadata); + mutableSegment.index(recordReader.next(reuse), METADATA); } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplAggregateMetricsTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplAggregateMetricsTest.java index 75b112c075d5..3dc4ed89d432 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplAggregateMetricsTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplAggregateMetricsTest.java @@ -36,6 +36,8 @@ import org.testng.Assert; import org.testng.annotations.Test; +import static org.mockito.Mockito.mock; + public class MutableSegmentImplAggregateMetricsTest { private static final String DIMENSION_1 = "dim1"; @@ -46,6 +48,7 @@ public class MutableSegmentImplAggregateMetricsTest { private static final String TIME_COLUMN2 = "time2"; private static final String KEY_SEPARATOR = "\t\t"; private static final int NUM_ROWS = 10001; + private static final StreamMessageMetadata METADATA = mock(StreamMessageMetadata.class); @Test public void testAggregateMetrics() @@ -99,7 +102,6 @@ private void testAggregateMetrics(MutableSegmentImpl mutableSegmentImpl) Map expectedValues = new HashMap<>(); Map expectedValuesFloat = new HashMap<>(); - StreamMessageMetadata defaultMetadata = new StreamMessageMetadata(System.currentTimeMillis(), new GenericRow()); for (int i = 0; i < NUM_ROWS; i++) { int hoursSinceEpoch = random.nextInt(10); int daysSinceEpoch = random.nextInt(5); @@ -114,7 +116,7 @@ private void testAggregateMetrics(MutableSegmentImpl mutableSegmentImpl) float metricValueFloat = floatValues[random.nextInt(floatValues.length)]; row.putValue(METRIC_2, metricValueFloat); - mutableSegmentImpl.index(row, defaultMetadata); + mutableSegmentImpl.index(row, METADATA); // Update expected values String key = buildKey(row); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplIngestionAggregationTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplIngestionAggregationTest.java index 455689cc6d45..eb8fbc30d142 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplIngestionAggregationTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplIngestionAggregationTest.java @@ -42,6 +42,8 @@ import org.testng.Assert; import org.testng.annotations.Test; +import static org.mockito.Mockito.mock; + public class MutableSegmentImplIngestionAggregationTest { private static final String DIMENSION_1 = "dim1"; @@ -55,6 +57,8 @@ public class MutableSegmentImplIngestionAggregationTest { private static final String KEY_SEPARATOR = "\t\t"; private static final int NUM_ROWS = 10001; + private static final StreamMessageMetadata METADATA = mock(StreamMessageMetadata.class); + private static Schema.SchemaBuilder getSchemaBuilder() { return new Schema.SchemaBuilder().setSchemaName("testSchema") .addSingleValueDimension(DIMENSION_1, FieldSpec.DataType.INT) @@ -150,13 +154,12 @@ public void testValuesAreNullThrowsException() long seed = 2; Random random = new Random(seed); - StreamMessageMetadata defaultMetadata = new StreamMessageMetadata(System.currentTimeMillis(), null); // Generate random int to prevent overflow GenericRow row = getRow(random, 1); row.putValue(METRIC, null); try { - mutableSegmentImpl.index(row, defaultMetadata); + mutableSegmentImpl.index(row, METADATA); Assert.fail(); } catch (NullPointerException e) { // expected @@ -311,7 +314,6 @@ private List addRowsDistinctCountHLL(long seed, MutableSegmentImpl mutab List metrics = new ArrayList<>(); Random random = new Random(seed); - StreamMessageMetadata defaultMetadata = new StreamMessageMetadata(System.currentTimeMillis(), null); HashMap hllMap = new HashMap<>(); HashMap> distinctMap = new HashMap<>(); @@ -335,7 +337,7 @@ private List addRowsDistinctCountHLL(long seed, MutableSegmentImpl mutab distinctMap.put(key, new HashSet<>(metricValue)); } - mutableSegmentImpl.index(row, defaultMetadata); + mutableSegmentImpl.index(row, METADATA); } distinctMap.forEach( @@ -355,7 +357,6 @@ private List addRowsSumPrecision(long seed, MutableSegmentImpl mutableSe List metrics = new ArrayList<>(); Random random = new Random(seed); - StreamMessageMetadata defaultMetadata = new StreamMessageMetadata(System.currentTimeMillis(), null); HashMap bdMap = new HashMap<>(); HashMap> bdIndividualMap = new HashMap<>(); @@ -378,7 +379,7 @@ private List addRowsSumPrecision(long seed, MutableSegmentImpl mutableSe bdIndividualMap.put(key, bdList); } - mutableSegmentImpl.index(row, defaultMetadata); + mutableSegmentImpl.index(row, METADATA); } for (String key : bdMap.keySet()) { @@ -400,7 +401,6 @@ private List> addRows(long seed, MutableSegmentImpl mutableSegmentI Set keys = new HashSet<>(); Random random = new Random(seed); - StreamMessageMetadata defaultMetadata = new StreamMessageMetadata(System.currentTimeMillis(), new GenericRow()); for (int i = 0; i < NUM_ROWS; i++) { // Generate random int to prevent overflow @@ -410,7 +410,7 @@ private List> addRows(long seed, MutableSegmentImpl mutableSegmentI row.putValue(METRIC, metricValue); row.putValue(METRIC_2, metric2Value); - mutableSegmentImpl.index(row, defaultMetadata); + mutableSegmentImpl.index(row, METADATA); String key = buildKey(row); metrics.add(Arrays.asList(new Metric(key, metricValue), new Metric(key, metric2Value))); @@ -465,7 +465,6 @@ public void testBigDecimalTooBig() { int seed = 1; Random random = new Random(seed); - StreamMessageMetadata defaultMetadata = new StreamMessageMetadata(System.currentTimeMillis(), null); MutableSegmentImpl mutableSegmentImpl = MutableSegmentImplTestUtils.createMutableSegmentImpl(schema, Collections.singleton(m1), VAR_LENGTH_SET, @@ -477,7 +476,7 @@ public void testBigDecimalTooBig() { row.putValue("metric", large); Assert.assertThrows(IllegalArgumentException.class, () -> { - mutableSegmentImpl.index(row, defaultMetadata); + mutableSegmentImpl.index(row, METADATA); }); mutableSegmentImpl.destroy(); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplRawMVTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplRawMVTest.java index b6da06bc63d5..cf33bcf092d6 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplRawMVTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplRawMVTest.java @@ -57,6 +57,8 @@ import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertNull; @@ -104,15 +106,16 @@ public void setUp() _mutableSegmentImpl = MutableSegmentImplTestUtils.createMutableSegmentImpl(_schema, new HashSet<>(noDictionaryColumns), Set.of(), Set.of(), false); - _lastIngestionTimeMs = System.currentTimeMillis(); - StreamMessageMetadata defaultMetadata = new StreamMessageMetadata(_lastIngestionTimeMs, new GenericRow()); - _startTimeMs = System.currentTimeMillis(); - + long currentTimeMs = System.currentTimeMillis(); + StreamMessageMetadata metadata = mock(StreamMessageMetadata.class); + when(metadata.getRecordIngestionTimeMs()).thenReturn(currentTimeMs); + _lastIngestionTimeMs = currentTimeMs; + _startTimeMs = currentTimeMs; try (RecordReader recordReader = RecordReaderFactory.getRecordReader(FileFormat.AVRO, avroFile, _schema.getColumnNames(), null)) { GenericRow reuse = new GenericRow(); while (recordReader.hasNext()) { - _mutableSegmentImpl.index(recordReader.next(reuse), defaultMetadata); + _mutableSegmentImpl.index(recordReader.next(reuse), metadata); _lastIndexedTs = System.currentTimeMillis(); } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplTest.java index 223033b29490..64a432d0390c 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImplTest.java @@ -49,6 +49,8 @@ import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; @@ -83,15 +85,16 @@ public void setUp() _schema = config.getSchema(); VirtualColumnProviderFactory.addBuiltInVirtualColumnsToSegmentSchema(_schema, "testSegment"); _mutableSegmentImpl = MutableSegmentImplTestUtils.createMutableSegmentImpl(_schema); - _lastIngestionTimeMs = System.currentTimeMillis(); - StreamMessageMetadata defaultMetadata = new StreamMessageMetadata(_lastIngestionTimeMs, new GenericRow()); - _startTimeMs = System.currentTimeMillis(); - - try (RecordReader recordReader = RecordReaderFactory - .getRecordReader(FileFormat.AVRO, avroFile, _schema.getColumnNames(), null)) { + long currentTimeMs = System.currentTimeMillis(); + StreamMessageMetadata metadata = mock(StreamMessageMetadata.class); + when(metadata.getRecordIngestionTimeMs()).thenReturn(currentTimeMs); + _lastIngestionTimeMs = currentTimeMs; + _startTimeMs = currentTimeMs; + try (RecordReader recordReader = RecordReaderFactory.getRecordReader(FileFormat.AVRO, avroFile, + _schema.getColumnNames(), null)) { GenericRow reuse = new GenericRow(); while (recordReader.hasNext()) { - _mutableSegmentImpl.index(recordReader.next(reuse), defaultMetadata); + _mutableSegmentImpl.index(recordReader.next(reuse), metadata); _lastIndexedTs = System.currentTimeMillis(); } } diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/MutableSegment.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/MutableSegment.java index 7c4d0e172917..910ec30f9b3b 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/MutableSegment.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/MutableSegment.java @@ -22,7 +22,7 @@ import java.io.IOException; import javax.annotation.Nullable; import org.apache.pinot.spi.data.readers.GenericRow; -import org.apache.pinot.spi.stream.RowMetadata; +import org.apache.pinot.spi.stream.StreamMessageMetadata; public interface MutableSegment extends IndexSegment { @@ -31,10 +31,10 @@ public interface MutableSegment extends IndexSegment { * Indexes a record into the segment with optionally provided metadata. * * @param row Record represented as a {@link GenericRow} - * @param rowMetadata the metadata associated with the message + * @param metadata the metadata associated with the message * @return Whether the segment is full (i.e. cannot index more record into it) */ - boolean index(GenericRow row, @Nullable RowMetadata rowMetadata) + boolean index(GenericRow row, @Nullable StreamMessageMetadata metadata) throws IOException; /** diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/stream/BytesStreamMessage.java b/pinot-spi/src/main/java/org/apache/pinot/spi/stream/BytesStreamMessage.java index 0b3ffb1fcace..7eda94bf4323 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/stream/BytesStreamMessage.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/stream/BytesStreamMessage.java @@ -23,15 +23,11 @@ public class BytesStreamMessage extends StreamMessage { - public BytesStreamMessage(@Nullable byte[] key, byte[] value, @Nullable StreamMessageMetadata metadata) { + public BytesStreamMessage(@Nullable byte[] key, byte[] value, StreamMessageMetadata metadata) { super(key, value, value.length, metadata); } - public BytesStreamMessage(byte[] value, @Nullable StreamMessageMetadata metadata) { + public BytesStreamMessage(byte[] value, StreamMessageMetadata metadata) { this(null, value, metadata); } - - public BytesStreamMessage(byte[] value) { - this(value, null); - } } diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/stream/ConsumerPartitionState.java b/pinot-spi/src/main/java/org/apache/pinot/spi/stream/ConsumerPartitionState.java index 668f7710772c..d752877f2a31 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/stream/ConsumerPartitionState.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/stream/ConsumerPartitionState.java @@ -29,10 +29,11 @@ public class ConsumerPartitionState { private final StreamPartitionMsgOffset _currentOffset; private final long _lastProcessedTimeMs; private final StreamPartitionMsgOffset _upstreamLatestOffset; - private final RowMetadata _lastProcessedRowMetadata; + private final StreamMessageMetadata _lastProcessedRowMetadata; public ConsumerPartitionState(String partitionId, StreamPartitionMsgOffset currentOffset, long lastProcessedTimeMs, - @Nullable StreamPartitionMsgOffset upstreamLatestOffset, @Nullable RowMetadata lastProcessedRowMetadata) { + @Nullable StreamPartitionMsgOffset upstreamLatestOffset, + @Nullable StreamMessageMetadata lastProcessedRowMetadata) { _partitionId = partitionId; _currentOffset = currentOffset; _lastProcessedTimeMs = lastProcessedTimeMs; @@ -56,7 +57,7 @@ public StreamPartitionMsgOffset getUpstreamLatestOffset() { return _upstreamLatestOffset; } - public RowMetadata getLastProcessedRowMetadata() { + public StreamMessageMetadata getLastProcessedRowMetadata() { return _lastProcessedRowMetadata; } } diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/stream/MessageBatch.java b/pinot-spi/src/main/java/org/apache/pinot/spi/stream/MessageBatch.java index 9a2eae1c3065..0af341e0d903 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/stream/MessageBatch.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/stream/MessageBatch.java @@ -47,19 +47,12 @@ default int getUnfilteredMessageCount() { /** * Returns the stream message at the given index within the batch. */ - default StreamMessage getStreamMessage(int index) { - byte[] value = getMessageBytesAtIndex(index); - StreamMessageMetadata metadata = (StreamMessageMetadata) getMetadataAtIndex(index); - //noinspection unchecked - return (StreamMessage) new StreamMessage<>(value, value.length, metadata); - } + StreamMessage getStreamMessage(int index); /** * Returns the start offset of the next batch. */ - default StreamPartitionMsgOffset getOffsetOfNextBatch() { - return getNextStreamPartitionMsgOffsetAtIndex(getMessageCount() - 1); - } + StreamPartitionMsgOffset getOffsetOfNextBatch(); /** * Returns the offset of the first message (including tombstone) in the batch. @@ -71,8 +64,7 @@ default StreamPartitionMsgOffset getFirstMessageOffset() { if (numMessages == 0) { return null; } - StreamMessageMetadata firstMessageMetadata = getStreamMessage(0).getMetadata(); - return firstMessageMetadata != null ? firstMessageMetadata.getOffset() : null; + return getStreamMessage(0).getMetadata().getOffset(); } /** @@ -103,39 +95,4 @@ default boolean isEndOfPartitionGroup() { default boolean hasDataLoss() { return false; } - - @Deprecated - default T getMessageAtIndex(int index) { - throw new UnsupportedOperationException(); - } - - @Deprecated - default byte[] getMessageBytesAtIndex(int index) { - return (byte[]) getMessageAtIndex(index); - } - - @Deprecated - default int getMessageLengthAtIndex(int index) { - throw new UnsupportedOperationException(); - } - - @Deprecated - default int getMessageOffsetAtIndex(int index) { - throw new UnsupportedOperationException(); - } - - @Deprecated - default RowMetadata getMetadataAtIndex(int index) { - return null; - } - - @Deprecated - default long getNextStreamMessageOffsetAtIndex(int index) { - throw new UnsupportedOperationException(); - } - - @Deprecated - default StreamPartitionMsgOffset getNextStreamPartitionMsgOffsetAtIndex(int index) { - return new LongMsgOffset(getNextStreamMessageOffsetAtIndex(index)); - } } diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/stream/RowMetadata.java b/pinot-spi/src/main/java/org/apache/pinot/spi/stream/RowMetadata.java deleted file mode 100644 index 459bc8cf9535..000000000000 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/stream/RowMetadata.java +++ /dev/null @@ -1,102 +0,0 @@ -/** - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ -package org.apache.pinot.spi.stream; - -import java.util.Map; -import javax.annotation.Nullable; -import org.apache.pinot.spi.annotations.InterfaceAudience; -import org.apache.pinot.spi.annotations.InterfaceStability; -import org.apache.pinot.spi.data.readers.GenericRow; - - -/** - * A class that provides relevant row-level metadata for rows indexed into a segment. - * - * Currently this is relevant for rows ingested into a mutable segment - the metadata is expected to be - * provided by the underlying stream. - */ -@InterfaceAudience.Public -@InterfaceStability.Evolving -public interface RowMetadata { - - /** - * Returns the timestamp associated with the record. This typically refers to the time it was ingested into the - * (last) upstream source. In some cases, it may be the time at which the record was created, aka event time - * (eg. in kafka, a topic may be configured to use record `CreateTime` instead of `LogAppendTime`). - * - * Expected to be used for stream-based sources. - * - * @return timestamp (epoch in milliseconds) when the row was ingested upstream - * Long.MIN_VALUE if not available - */ - long getRecordIngestionTimeMs(); - - /** - * When supported by the underlying stream, this method returns the timestamp in milliseconds associated with - * the ingestion of the record in the first stream. - * - * Complex ingestion pipelines may be composed of multiple streams: - * (EventCreation) -> {First Stream} -> ... -> {Last Stream} - * - * @return timestamp (epoch in milliseconds) when the row was initially ingested upstream for the first - * time Long.MIN_VALUE if not supported by the underlying stream. - */ - default long getFirstStreamRecordIngestionTimeMs() { - return Long.MIN_VALUE; - } - - /** - * @return The serialized size of the record - */ - default int getRecordSerializedSize() { - return Integer.MIN_VALUE; - } - - /** - * Returns the stream offset of the message. - */ - @Nullable - default StreamPartitionMsgOffset getOffset() { - return null; - } - - /** - * Returns the next stream offset of the message. - */ - @Nullable - default StreamPartitionMsgOffset getNextOffset() { - return null; - } - - /** - * Returns the stream message headers. - */ - @Nullable - default GenericRow getHeaders() { - return null; - } - - /** - * Returns the metadata associated with the stream message. - */ - @Nullable - default Map getRecordMetadata() { - return null; - } -} diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/stream/StreamDataDecoderImpl.java b/pinot-spi/src/main/java/org/apache/pinot/spi/stream/StreamDataDecoderImpl.java index 6bd31ca72cf4..96073ace839c 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/stream/StreamDataDecoderImpl.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/stream/StreamDataDecoderImpl.java @@ -61,17 +61,15 @@ public StreamDataDecoderResult decode(StreamMessage message) { row.putValue(KEY, new String(message.getKey(), StandardCharsets.UTF_8)); } StreamMessageMetadata metadata = message.getMetadata(); - if (metadata != null) { - GenericRow headers = metadata.getHeaders(); - if (headers != null) { - headers.getFieldToValueMap().forEach((k, v) -> row.putValue(HEADER_KEY_PREFIX + k, v)); - } - Map recordMetadata = metadata.getRecordMetadata(); - if (recordMetadata != null) { - recordMetadata.forEach((k, v) -> row.putValue(METADATA_KEY_PREFIX + k, v)); - } - row.putValue(RECORD_SERIALIZED_VALUE_SIZE_KEY, length); + GenericRow headers = metadata.getHeaders(); + if (headers != null) { + headers.getFieldToValueMap().forEach((k, v) -> row.putValue(HEADER_KEY_PREFIX + k, v)); } + Map recordMetadata = metadata.getRecordMetadata(); + if (recordMetadata != null) { + recordMetadata.forEach((k, v) -> row.putValue(METADATA_KEY_PREFIX + k, v)); + } + row.putValue(RECORD_SERIALIZED_VALUE_SIZE_KEY, length); return new StreamDataDecoderResult(row, null); } else { return new StreamDataDecoderResult(null, diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/stream/StreamMessage.java b/pinot-spi/src/main/java/org/apache/pinot/spi/stream/StreamMessage.java index e973e92a110b..96d27e48d599 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/stream/StreamMessage.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/stream/StreamMessage.java @@ -26,7 +26,7 @@ * 1. record key (optional) * 2. record value (required) * 3. length of the record value (required) - * 4. StreamMessageMetadata (optional) - encapsulates record headers and metadata associated with a stream message + * 4. StreamMessageMetadata (required) - encapsulates record headers and metadata associated with a stream message * (such as a message identifier, publish timestamp, user-provided headers etc) * * Similar to value decoder, each implementing stream plugin can have a key decoder and header extractor. @@ -39,35 +39,21 @@ * Additionally, the pinot table schema should refer these fields. Otherwise, even though the fields are extracted, * they will not materialize in the pinot table. */ +// TODO: Revisit if we need to support value type other than byte[] public class StreamMessage { protected final byte[] _key; protected final T _value; protected final int _length; protected final StreamMessageMetadata _metadata; - public StreamMessage(@Nullable byte[] key, T value, int length, @Nullable StreamMessageMetadata metadata) { + public StreamMessage(@Nullable byte[] key, T value, int length, StreamMessageMetadata metadata) { + assert value != null && metadata != null; _key = key; _value = value; _length = length; _metadata = metadata; } - public StreamMessage(T value, int length, @Nullable StreamMessageMetadata metadata) { - this(null, value, length, metadata); - } - - public StreamMessage(T value, int length) { - this(value, length, null); - } - - @Deprecated - public StreamMessage(@Nullable byte[] key, T value, @Nullable StreamMessageMetadata metadata, int length) { - _key = key; - _value = value; - _metadata = metadata; - _length = length; - } - /** * Returns the key of the message. */ @@ -93,7 +79,6 @@ public int getLength() { /** * Returns the metadata of the message. */ - @Nullable public StreamMessageMetadata getMetadata() { return _metadata; } diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/stream/StreamMessageMetadata.java b/pinot-spi/src/main/java/org/apache/pinot/spi/stream/StreamMessageMetadata.java index 152d109fbd60..0c9b38f07c69 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/stream/StreamMessageMetadata.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/stream/StreamMessageMetadata.java @@ -27,7 +27,7 @@ * A class that provides metadata associated with the message of a stream, for e.g., * timestamp derived from the incoming record (not the ingestion time). */ -public class StreamMessageMetadata implements RowMetadata { +public class StreamMessageMetadata { private final long _recordIngestionTimeMs; private final long _firstStreamRecordIngestionTimeMs; private final int _recordSerializedSize; @@ -36,30 +36,9 @@ public class StreamMessageMetadata implements RowMetadata { private final GenericRow _headers; private final Map _metadata; - @Deprecated - public StreamMessageMetadata(long recordIngestionTimeMs) { - this(recordIngestionTimeMs, null); - } - - @Deprecated - public StreamMessageMetadata(long recordIngestionTimeMs, @Nullable GenericRow headers) { - this(recordIngestionTimeMs, headers, null); - } - - @Deprecated - public StreamMessageMetadata(long recordIngestionTimeMs, @Nullable GenericRow headers, Map metadata) { - this(recordIngestionTimeMs, Long.MIN_VALUE, headers, metadata); - } - - @Deprecated - public StreamMessageMetadata(long recordIngestionTimeMs, long firstStreamRecordIngestionTimeMs, - @Nullable GenericRow headers, Map metadata) { - this(recordIngestionTimeMs, firstStreamRecordIngestionTimeMs, null, null, Integer.MIN_VALUE, headers, metadata); - } - - public StreamMessageMetadata(long recordIngestionTimeMs, long firstStreamRecordIngestionTimeMs, - @Nullable StreamPartitionMsgOffset offset, @Nullable StreamPartitionMsgOffset nextOffset, - int recordSerializedSize, @Nullable GenericRow headers, @Nullable Map metadata) { + private StreamMessageMetadata(long recordIngestionTimeMs, long firstStreamRecordIngestionTimeMs, + StreamPartitionMsgOffset offset, StreamPartitionMsgOffset nextOffset, int recordSerializedSize, + @Nullable GenericRow headers, @Nullable Map metadata) { _recordIngestionTimeMs = recordIngestionTimeMs; _firstStreamRecordIngestionTimeMs = firstStreamRecordIngestionTimeMs; _offset = offset; @@ -69,51 +48,42 @@ public StreamMessageMetadata(long recordIngestionTimeMs, long firstStreamRecordI _metadata = metadata; } - @Override public long getRecordIngestionTimeMs() { return _recordIngestionTimeMs; } - @Override public long getFirstStreamRecordIngestionTimeMs() { return _firstStreamRecordIngestionTimeMs; } - @Override public int getRecordSerializedSize() { return _recordSerializedSize; } - @Nullable - @Override public StreamPartitionMsgOffset getOffset() { return _offset; } - @Nullable - @Override public StreamPartitionMsgOffset getNextOffset() { return _nextOffset; } @Nullable - @Override public GenericRow getHeaders() { return _headers; } @Nullable - @Override public Map getRecordMetadata() { return _metadata; } public static class Builder { private long _recordIngestionTimeMs = Long.MIN_VALUE; - private int _recordSerializedSize = Integer.MIN_VALUE; private long _firstStreamRecordIngestionTimeMs = Long.MIN_VALUE; private StreamPartitionMsgOffset _offset; private StreamPartitionMsgOffset _nextOffset; + private int _recordSerializedSize = Integer.MIN_VALUE; private GenericRow _headers; private Map _metadata; @@ -133,6 +103,11 @@ public Builder setOffset(StreamPartitionMsgOffset offset, StreamPartitionMsgOffs return this; } + public Builder setSerializedValueSize(int recordSerializedSize) { + _recordSerializedSize = recordSerializedSize; + return this; + } + public Builder setHeaders(GenericRow headers) { _headers = headers; return this; @@ -143,12 +118,8 @@ public Builder setMetadata(Map metadata) { return this; } - public Builder setSerializedValueSize(int recordSerializedSize) { - _recordSerializedSize = recordSerializedSize; - return this; - } - public StreamMessageMetadata build() { + assert _offset != null && _nextOffset != null; return new StreamMessageMetadata(_recordIngestionTimeMs, _firstStreamRecordIngestionTimeMs, _offset, _nextOffset, _recordSerializedSize, _headers, _metadata); } diff --git a/pinot-spi/src/test/java/org/apache/pinot/spi/stream/StreamDataDecoderImplTest.java b/pinot-spi/src/test/java/org/apache/pinot/spi/stream/StreamDataDecoderImplTest.java index 4a8420ea5d65..0096691cc139 100644 --- a/pinot-spi/src/test/java/org/apache/pinot/spi/stream/StreamDataDecoderImplTest.java +++ b/pinot-spi/src/test/java/org/apache/pinot/spi/stream/StreamDataDecoderImplTest.java @@ -18,9 +18,7 @@ */ package org.apache.pinot.spi.stream; -import com.google.common.collect.ImmutableSet; import java.nio.charset.StandardCharsets; -import java.util.Collections; import java.util.Map; import java.util.Set; import javax.annotation.Nullable; @@ -28,38 +26,45 @@ import org.testng.Assert; import org.testng.annotations.Test; +import static org.mockito.Mockito.mock; + public class StreamDataDecoderImplTest { private static final String NAME_FIELD = "name"; private static final String AGE_HEADER_KEY = "age"; - private static final String SEQNO_RECORD_METADATA = "seqNo"; + private static final String SEQ_NO_RECORD_METADATA = "seqNo"; + private static final StreamMessageMetadata METADATA = mock(StreamMessageMetadata.class); @Test public void testDecodeValueOnly() { TestDecoder messageDecoder = new TestDecoder(); - messageDecoder.init(Collections.emptyMap(), ImmutableSet.of(NAME_FIELD), ""); + messageDecoder.init(Map.of(), Set.of(NAME_FIELD), ""); String value = "Alice"; - BytesStreamMessage message = new BytesStreamMessage(value.getBytes(StandardCharsets.UTF_8)); + BytesStreamMessage message = new BytesStreamMessage(value.getBytes(StandardCharsets.UTF_8), METADATA); StreamDataDecoderResult result = new StreamDataDecoderImpl(messageDecoder).decode(message); Assert.assertNotNull(result); Assert.assertNull(result.getException()); Assert.assertNotNull(result.getResult()); GenericRow row = result.getResult(); - Assert.assertEquals(row.getFieldToValueMap().size(), 1); + Assert.assertEquals(row.getFieldToValueMap().size(), 2); Assert.assertEquals(String.valueOf(row.getValue(NAME_FIELD)), value); + Assert.assertEquals(row.getValue(StreamDataDecoderImpl.RECORD_SERIALIZED_VALUE_SIZE_KEY), value.length()); } @Test public void testDecodeKeyAndHeaders() { TestDecoder messageDecoder = new TestDecoder(); - messageDecoder.init(Collections.emptyMap(), ImmutableSet.of(NAME_FIELD), ""); + messageDecoder.init(Map.of(), Set.of(NAME_FIELD), ""); String value = "Alice"; String key = "id-1"; GenericRow headers = new GenericRow(); headers.putValue(AGE_HEADER_KEY, 3); - Map recordMetadata = Collections.singletonMap(SEQNO_RECORD_METADATA, "1"); - StreamMessageMetadata metadata = new StreamMessageMetadata(1234L, headers, recordMetadata); + StreamMessageMetadata metadata = new StreamMessageMetadata.Builder().setRecordIngestionTimeMs(1234L) + .setOffset(new LongMsgOffset(0), new LongMsgOffset(1)) + .setHeaders(headers) + .setMetadata(Map.of(SEQ_NO_RECORD_METADATA, "1")) + .build(); BytesStreamMessage message = new BytesStreamMessage(key.getBytes(StandardCharsets.UTF_8), value.getBytes(StandardCharsets.UTF_8), metadata); @@ -73,16 +78,16 @@ public void testDecodeKeyAndHeaders() { Assert.assertEquals(row.getValue(NAME_FIELD), value); Assert.assertEquals(row.getValue(StreamDataDecoderImpl.KEY), key, "Failed to decode record key"); Assert.assertEquals(row.getValue(StreamDataDecoderImpl.HEADER_KEY_PREFIX + AGE_HEADER_KEY), 3); - Assert.assertEquals(row.getValue(StreamDataDecoderImpl.METADATA_KEY_PREFIX + SEQNO_RECORD_METADATA), "1"); + Assert.assertEquals(row.getValue(StreamDataDecoderImpl.METADATA_KEY_PREFIX + SEQ_NO_RECORD_METADATA), "1"); Assert.assertEquals(row.getValue(StreamDataDecoderImpl.RECORD_SERIALIZED_VALUE_SIZE_KEY), value.length()); } @Test public void testNoExceptionIsThrown() { ThrowingDecoder messageDecoder = new ThrowingDecoder(); - messageDecoder.init(Collections.emptyMap(), ImmutableSet.of(NAME_FIELD), ""); + messageDecoder.init(Map.of(), Set.of(NAME_FIELD), ""); String value = "Alice"; - BytesStreamMessage message = new BytesStreamMessage(value.getBytes(StandardCharsets.UTF_8)); + BytesStreamMessage message = new BytesStreamMessage(value.getBytes(StandardCharsets.UTF_8), METADATA); StreamDataDecoderResult result = new StreamDataDecoderImpl(messageDecoder).decode(message); Assert.assertNotNull(result); Assert.assertNotNull(result.getException()); From c1fceb23986f9b6aa460b2f307b72f811065f6c2 Mon Sep 17 00:00:00 2001 From: Rajat Venkatesh <1638298+vrajat@users.noreply.github.com> Date: Fri, 25 Jul 2025 05:41:32 +0530 Subject: [PATCH 022/167] Queries now self terminate if in panic mode. (#16380) * Queries now self terminate if in panic mode. * Add config test * Hard kill on critical level. * Fix configs * Separate anchor thread interruption. * Checkstyle * Review comments * remove code for critical level --------- Co-authored-by: Rajat Venkatesh --- .../PerQueryCPUMemAccountantFactory.java | 18 ++++++++++++++++ .../core/accounting/QueryMonitorConfig.java | 21 +++++++++++++++++++ .../accounting/QueryMonitorConfigTest.java | 17 +++++++++++++++ .../ThreadResourceUsageAccountant.java | 9 ++++++++ .../org/apache/pinot/spi/trace/Tracing.java | 3 ++- .../pinot/spi/utils/CommonConstants.java | 4 ++++ 6 files changed, 71 insertions(+), 1 deletion(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java index e2f23664e392..46e700fe8438 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java @@ -316,6 +316,16 @@ public boolean isAnchorThreadInterrupted() { return false; } + @Override + public boolean isQueryTerminated() { + QueryMonitorConfig config = _watcherTask.getQueryMonitorConfig(); + if (config.isThreadSelfTerminate() && _watcherTask.getHeapUsageBytes() > config.getPanicLevel()) { + logSelfTerminatedQuery(_threadLocalEntry.get().getQueryId(), Thread.currentThread()); + return true; + } + return false; + } + @Override @Deprecated public void createExecutionContext(String queryId, int taskId, ThreadExecutionContext.TaskType taskType, @@ -581,6 +591,14 @@ protected void logTerminatedQuery(QueryResourceTracker queryResourceTracker, lon queryResourceTracker.getCpuTimeNs(), totalHeapMemoryUsage, hasCallback); } + protected void logSelfTerminatedQuery(String queryId, Thread queryThread) { + if (!_cancelSentQueries.contains(queryId)) { + LOGGER.warn("{} self-terminated. Heap Usage: {}. Query Thread: {}", + queryId, _watcherTask.getHeapUsageBytes(), queryThread.getName()); + _cancelSentQueries.add(queryId); + } + } + @Override public Exception getErrorStatus() { return _threadLocalEntry.get()._errorStatus.getAndSet(null); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/accounting/QueryMonitorConfig.java b/pinot-core/src/main/java/org/apache/pinot/core/accounting/QueryMonitorConfig.java index 327b34e4bd2c..3efbc432e636 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/accounting/QueryMonitorConfig.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/accounting/QueryMonitorConfig.java @@ -62,6 +62,8 @@ public class QueryMonitorConfig { private final boolean _isQueryKilledMetricEnabled; + private final boolean _isThreadSelfTerminate; + public QueryMonitorConfig(PinotConfiguration config, long maxHeapSize) { _maxHeapSize = maxHeapSize; @@ -106,6 +108,9 @@ public QueryMonitorConfig(PinotConfiguration config, long maxHeapSize) { _isQueryKilledMetricEnabled = config.getProperty(CommonConstants.Accounting.CONFIG_OF_QUERY_KILLED_METRIC_ENABLED, CommonConstants.Accounting.DEFAULT_QUERY_KILLED_METRIC_ENABLED); + + _isThreadSelfTerminate = config.getProperty(CommonConstants.Accounting.CONFIG_OF_THREAD_SELF_TERMINATE, + CommonConstants.Accounting.DEFAULT_THREAD_SELF_TERMINATE); } QueryMonitorConfig(QueryMonitorConfig oldConfig, Set changedConfigs, Map clusterConfigs) { @@ -245,6 +250,18 @@ public QueryMonitorConfig(PinotConfiguration config, long maxHeapSize) { } else { _isQueryKilledMetricEnabled = oldConfig._isQueryKilledMetricEnabled; } + + if (changedConfigs.contains(CommonConstants.Accounting.CONFIG_OF_THREAD_SELF_TERMINATE)) { + if (clusterConfigs == null || !clusterConfigs.containsKey( + CommonConstants.Accounting.CONFIG_OF_THREAD_SELF_TERMINATE)) { + _isThreadSelfTerminate = CommonConstants.Accounting.DEFAULT_THREAD_SELF_TERMINATE; + } else { + _isThreadSelfTerminate = + Boolean.parseBoolean(clusterConfigs.get(CommonConstants.Accounting.CONFIG_OF_THREAD_SELF_TERMINATE)); + } + } else { + _isThreadSelfTerminate = oldConfig._isThreadSelfTerminate; + } } public long getMaxHeapSize() { @@ -294,4 +311,8 @@ public long getCpuTimeBasedKillingThresholdNS() { public boolean isQueryKilledMetricEnabled() { return _isQueryKilledMetricEnabled; } + + public boolean isThreadSelfTerminate() { + return _isThreadSelfTerminate; + } } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/accounting/QueryMonitorConfigTest.java b/pinot-core/src/test/java/org/apache/pinot/core/accounting/QueryMonitorConfigTest.java index d5afcd63d661..79cd17c6e8f8 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/accounting/QueryMonitorConfigTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/accounting/QueryMonitorConfigTest.java @@ -44,6 +44,7 @@ public class QueryMonitorConfigTest { private static final boolean EXPECTED_IS_CPU_TIME_BASED_KILLING_ENABLED = true; private static final long EXPECTED_CPU_TIME_BASED_KILLING_THRESHOLD_NS = 1000; private static final boolean EXPECTED_IS_QUERY_KILLED_METRIC_ENABLED = true; + private static final boolean EXPECTED_IS_THREAD_SELF_TERMINATE_IN_PANIC_MODE = true; private static final Map CLUSTER_CONFIGS = new HashMap<>(); private static String getFullyQualifiedConfigName(String config) { @@ -79,6 +80,9 @@ public void setUp() { Double.toString(EXPECTED_MIN_MEMORY_FOOTPRINT_FOR_KILL)); CLUSTER_CONFIGS.put(getFullyQualifiedConfigName(CommonConstants.Accounting.CONFIG_OF_QUERY_KILLED_METRIC_ENABLED), Boolean.toString(EXPECTED_IS_QUERY_KILLED_METRIC_ENABLED)); + CLUSTER_CONFIGS.put( + getFullyQualifiedConfigName(CommonConstants.Accounting.CONFIG_OF_THREAD_SELF_TERMINATE), + Boolean.toString(EXPECTED_IS_THREAD_SELF_TERMINATE_IN_PANIC_MODE)); } @Test @@ -243,4 +247,17 @@ void testQueryKilledMetricEnabledConfigChange() { CLUSTER_CONFIGS); assertTrue(accountant.getWatcherTask().getQueryMonitorConfig().isQueryKilledMetricEnabled()); } + + @Test + void testThreadSelfTerminateInPanicMode() { + PerQueryCPUMemAccountantFactory.PerQueryCPUMemResourceUsageAccountant accountant = + new PerQueryCPUMemAccountantFactory.PerQueryCPUMemResourceUsageAccountant(new PinotConfiguration(), "test", + InstanceType.SERVER); + + assertFalse(accountant.getWatcherTask().getQueryMonitorConfig().isThreadSelfTerminate()); + accountant.getWatcherTask().onChange( + Set.of(getFullyQualifiedConfigName(CommonConstants.Accounting.CONFIG_OF_THREAD_SELF_TERMINATE)), + CLUSTER_CONFIGS); + assertTrue(accountant.getWatcherTask().getQueryMonitorConfig().isThreadSelfTerminate()); + } } diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceUsageAccountant.java b/pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceUsageAccountant.java index 4fdb4f1e1b60..dc912371980f 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceUsageAccountant.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceUsageAccountant.java @@ -37,6 +37,15 @@ public interface ThreadResourceUsageAccountant { */ boolean isAnchorThreadInterrupted(); + /** + * This function is expected to be called by threads in query engine. The query id of the task will be available in + * the thread execution context stored in a thread local. Therefore it does not accept any parameters. + * @return true if the query is terminated, false otherwise + */ + default boolean isQueryTerminated() { + return false; + } + /** * This method has been deprecated and replaced by {@link setupRunner(String, int, ThreadExecutionContext.TaskType)} * and {@link setupWorker(int, ThreadExecutionContext.TaskType, ThreadExecutionContext)}. diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/trace/Tracing.java b/pinot-spi/src/main/java/org/apache/pinot/spi/trace/Tracing.java index afc9a283c1b8..831e6d08b6b7 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/trace/Tracing.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/trace/Tracing.java @@ -380,7 +380,8 @@ public static void startThreadAccountant() { } public static boolean isInterrupted() { - return Thread.interrupted() || Tracing.getThreadAccountant().isAnchorThreadInterrupted(); + return Thread.interrupted() || Tracing.getThreadAccountant().isAnchorThreadInterrupted() + || Tracing.getThreadAccountant().isQueryTerminated(); } public static void sampleAndCheckInterruption() { diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java index 306f8d710ed9..4a7676667e56 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java @@ -1548,6 +1548,10 @@ public static class Accounting { "accounting.cancel.callback.cache.expiry.seconds"; public static final int DEFAULT_CANCEL_CALLBACK_CACHE_EXPIRY_SECONDS = 1200; + public static final String CONFIG_OF_THREAD_SELF_TERMINATE = + "accounting.thread.self.terminate"; + public static final boolean DEFAULT_THREAD_SELF_TERMINATE = false; + /** * QUERY WORKLOAD ISOLATION Configs * From 4f1c40c6f1d84a441cf5998f7b4c3ab0f00880ac Mon Sep 17 00:00:00 2001 From: avshenuk <4578760+avshenuk@users.noreply.github.com> Date: Fri, 25 Jul 2025 02:17:57 +0200 Subject: [PATCH 023/167] feat: Enhance PurgeTask to automatically delete empty segments (#16368) --- .../PurgeMinionClusterIntegrationTest.java | 178 +++++++++++++++++- .../tasks/purge/PurgeTaskGenerator.java | 48 ++++- 2 files changed, 221 insertions(+), 5 deletions(-) diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PurgeMinionClusterIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PurgeMinionClusterIntegrationTest.java index 19459894c722..9aab0fc5771e 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PurgeMinionClusterIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PurgeMinionClusterIntegrationTest.java @@ -61,6 +61,8 @@ public class PurgeMinionClusterIntegrationTest extends BaseClusterIntegrationTes private static final String PURGE_DELTA_PASSED_TABLE = "myTable2"; private static final String PURGE_DELTA_NOT_PASSED_TABLE = "myTable3"; private static final String PURGE_OLD_SEGMENTS_WITH_NEW_INDICES_TABLE = "myTable4"; + private static final String PURGE_ALL_RECORDS_TABLE = "myTable5"; + private static final String PURGE_REALTIME_LAST_SEGMENT_TABLE = "myTable6"; protected PinotHelixTaskResourceManager _helixTaskResourceManager; protected PinotTaskManager _taskManager; @@ -82,11 +84,15 @@ public void setUp() startServer(); startMinion(); - List allTables = List.of(PURGE_FIRST_RUN_TABLE, PURGE_DELTA_PASSED_TABLE, PURGE_DELTA_NOT_PASSED_TABLE, - PURGE_OLD_SEGMENTS_WITH_NEW_INDICES_TABLE); + // Start Kafka for realtime table test + startKafka(); + + List allOfflineTables = + List.of(PURGE_FIRST_RUN_TABLE, PURGE_DELTA_PASSED_TABLE, PURGE_DELTA_NOT_PASSED_TABLE, + PURGE_OLD_SEGMENTS_WITH_NEW_INDICES_TABLE, PURGE_ALL_RECORDS_TABLE); Schema schema = null; TableConfig tableConfig = null; - for (String tableName : allTables) { + for (String tableName : allOfflineTables) { // create and upload schema schema = createSchema(); schema.setSchemaName(tableName); @@ -107,10 +113,24 @@ public void setUp() _segmentTarDir); // Upload segments for all tables - for (String tableName : allTables) { + for (String tableName : allOfflineTables) { uploadSegments(tableName, _segmentTarDir); } + // Set up realtime table with purge task configuration + schema = createSchema(); + schema.setSchemaName(PURGE_REALTIME_LAST_SEGMENT_TABLE); + addSchema(schema); + // Create realtime table config with purge task + TableConfig realtimeTableConfig = createRealtimeTableConfig(avroFiles.get(0)); + realtimeTableConfig.setTableName(PURGE_REALTIME_LAST_SEGMENT_TABLE); + realtimeTableConfig.setTaskConfig(getPurgeTaskConfig()); + addTableConfig(realtimeTableConfig); + // Push data into Kafka to create LLC segments + pushAvroIntoKafka(avroFiles); + // Wait for all documents loaded + waitForDocsLoaded(600_000L, true, PURGE_REALTIME_LAST_SEGMENT_TABLE); + setRecordPurger(); _helixTaskResourceManager = _controllerStarter.getHelixTaskResourceManager(); _taskManager = _controllerStarter.getTaskManager(); @@ -151,6 +171,10 @@ private void setRecordPurger() { PURGE_OLD_SEGMENTS_WITH_NEW_INDICES_TABLE); if (tableNames.contains(rawTableName)) { return row -> row.getValue("ArrTime").equals(1); + } else if (PURGE_ALL_RECORDS_TABLE.equals(rawTableName) || PURGE_REALTIME_LAST_SEGMENT_TABLE.equals( + rawTableName)) { + // Purge ALL records to test segment deletion + return row -> true; } else { return null; } @@ -381,6 +405,151 @@ public void testPurgeOnOldSegmentsWithIndicesOnNewColumns() verifyTableDelete(offlineTableName); } + /** + * Test that segments are automatically deleted when all records are purged + */ + @Test + public void testSegmentDeletionWhenAllRecordsPurged() + throws Exception { + // Expected behavior: + // 1. First run: All records in segments are purged (RecordPurger returns true for all records) + // 2. First run: Segments become empty but are still present with totalDocs = 0 + // 3. Second run: Empty segments are automatically deleted by PurgeTaskGenerator during task generation + // 4. Verify that segments are removed from the table + + String offlineTableName = TableNameBuilder.OFFLINE.tableNameWithType(PURGE_ALL_RECORDS_TABLE); + + // Get initial segment count + List initialSegments = _pinotHelixResourceManager.getSegmentsZKMetadata(offlineTableName); + int initialSegmentCount = initialSegments.size(); + assertTrue(initialSegmentCount > 0, "Table should have segments initially"); + + // First run: Schedule purge task to create empty segments + assertNotNull(_taskManager.scheduleTasks(new TaskSchedulingContext() + .setTablesToSchedule(Collections.singleton(offlineTableName))) + .get(MinionConstants.PurgeTask.TASK_TYPE)); + assertTrue(_helixTaskResourceManager.getTaskQueues() + .contains(PinotHelixTaskResourceManager.getHelixJobQueueName(MinionConstants.PurgeTask.TASK_TYPE))); + + // Wait for first task to complete + waitForTaskToComplete(); + + // Verify table now has no data but segments still exist (empty segments) + TestUtils.waitForCondition(aVoid -> getCurrentCountStarResult(PURGE_ALL_RECORDS_TABLE) == 0, 60_000L, + "Failed to get expected purged records"); + List segmentsAfterFirstRun = _pinotHelixResourceManager.getSegmentsZKMetadata(offlineTableName); + assertEquals(segmentsAfterFirstRun.size(), initialSegmentCount, + "Segments should still exist after first purge run"); + + // Verify segments have totalDocs = 0 + for (SegmentZKMetadata segment : segmentsAfterFirstRun) { + assertEquals(segment.getTotalDocs(), 0L, "All segments should have zero documents after purging"); + } + + // Second run: Schedule purge task again - this should delete the empty segments during task generation + assertNotNull(_taskManager.scheduleTasks(new TaskSchedulingContext() + .setTablesToSchedule(Collections.singleton(offlineTableName))) + .get(MinionConstants.PurgeTask.TASK_TYPE)); + + // Wait for second task to complete (if any tasks were generated) + waitForTaskToComplete(); + + // Verify that all empty segments have been deleted + TestUtils.waitForCondition(aVoid -> { + List remainingSegments = _pinotHelixResourceManager.getSegmentsZKMetadata(offlineTableName); + return remainingSegments.isEmpty(); + }, 60_000L, "Expected all empty segments to be deleted after second purge run"); + + // Verify table still has no data + assertEquals(getCurrentCountStarResult(PURGE_ALL_RECORDS_TABLE), 0); + + // Drop the table + dropOfflineTable(PURGE_ALL_RECORDS_TABLE); + + // Check if the task metadata is cleaned up on table deletion + verifyTableDelete(offlineTableName); + } + + /** + * Test that empty segments are preserved when they are the last segment of a partition in realtime tables. + * This test specifically covers the edge case where empty segments should only + * be allowed when they are needed to mark the end of a stream partition (e.g. Kinesis). + */ + @Test + public void testRealtimeLastSegmentPreservation() + throws Exception { + // Expected behavior: + // 1. First run: All records in completed segments are purged (RecordPurger returns true for all records) + // 2. First run: Completed segments become empty but are still present with totalDocs = 0 + // 3. Second run: Empty non-last completed segments are deleted, but last segments per partition are preserved + // 4. Verify that consuming segments and last empty completed segments per partition remain + + String realtimeTableName = TableNameBuilder.REALTIME.tableNameWithType(PURGE_REALTIME_LAST_SEGMENT_TABLE); + + // Get initial segment count + List initialSegments = _pinotHelixResourceManager.getSegmentsZKMetadata(realtimeTableName); + int initialSegmentCount = initialSegments.size(); + assertTrue(initialSegmentCount > 0, "Table should have segments initially"); + + // First run: Schedule purge task to create empty segments + assertNotNull(_taskManager.scheduleTasks(new TaskSchedulingContext() + .setTablesToSchedule(Collections.singleton(realtimeTableName))) + .get(MinionConstants.PurgeTask.TASK_TYPE)); + + // Wait for first task to complete + waitForTaskToComplete(); + + // Calculate expected remaining records after purging completed segments + // Expected remaining = totalRecords % recordsPerSegmentPerPartition + long expectedRemainingRecords = getCountStarResult() % (getRealtimeSegmentFlushSize() / getNumKafkaPartitions()); + TestUtils.waitForCondition( + aVoid -> getCurrentCountStarResult(PURGE_REALTIME_LAST_SEGMENT_TABLE) == expectedRemainingRecords, + 60_000L, "Failed to get expected purged records"); + List segmentsAfterFirstRun = _pinotHelixResourceManager.getSegmentsZKMetadata(realtimeTableName); + assertEquals(segmentsAfterFirstRun.size(), initialSegmentCount, + "Segments should still exist after first purge run"); + + // Verify segments have totalDocs = 0 (for completed segments) + for (SegmentZKMetadata segment : segmentsAfterFirstRun) { + if (segment.getStatus().isCompleted()) { + assertEquals(segment.getTotalDocs(), 0L, "All completed segments should have zero documents after purging"); + } + } + + // Second run: Schedule purge task again - this should delete empty non-last segments but preserve last segments + assertNotNull(_taskManager.scheduleTasks(new TaskSchedulingContext() + .setTablesToSchedule(Collections.singleton(realtimeTableName))) + .get(MinionConstants.PurgeTask.TASK_TYPE)); + + // Wait for second task to complete + waitForTaskToComplete(); + + TestUtils.waitForCondition(aVoid -> { + // Verify that we have the expected number of segments remaining: + // - 1 consuming segment per partition (should not be deleted) + // - 1 completed empty segment per partition (last segment, should be preserved) + List remainingSegments = _pinotHelixResourceManager.getSegmentsZKMetadata(realtimeTableName); + long consumingSegments = remainingSegments.stream() + .filter(s -> !s.getStatus().isCompleted()) + .count(); + long completedSegments = remainingSegments.stream() + .filter(s -> s.getStatus().isCompleted()) + .count(); + + // We should have: 1 consuming segment per partition + 1 last empty completed segment per partition + return consumingSegments == getNumKafkaPartitions() && completedSegments == getNumKafkaPartitions(); + }, 60_000L, "Expected all but last empty completed segments to be deleted after second purge run"); + + // Verify table still has expected remaining data (from consuming segments) + assertEquals(getCurrentCountStarResult(PURGE_REALTIME_LAST_SEGMENT_TABLE), expectedRemainingRecords); + + // Drop the realtime table + dropRealtimeTable(PURGE_REALTIME_LAST_SEGMENT_TABLE); + + // Verify cleanup + verifyTableDelete(realtimeTableName); + } + protected void verifyTableDelete(String tableNameWithType) { TestUtils.waitForCondition(input -> { // Check if the segment lineage is cleaned up @@ -416,6 +585,7 @@ public void tearDown() stopServer(); stopBroker(); stopController(); + stopKafka(); stopZk(); FileUtils.deleteDirectory(_tempDir); } diff --git a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/purge/PurgeTaskGenerator.java b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/purge/PurgeTaskGenerator.java index 3e0200959740..bc0a6f1098f4 100644 --- a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/purge/PurgeTaskGenerator.java +++ b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/purge/PurgeTaskGenerator.java @@ -23,11 +23,13 @@ import java.util.Collections; import java.util.Comparator; import java.util.HashMap; +import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Set; import org.apache.pinot.common.data.Segment; import org.apache.pinot.common.metadata.segment.SegmentZKMetadata; +import org.apache.pinot.common.utils.LLCSegmentName; import org.apache.pinot.controller.helix.core.minion.generator.BaseTaskGenerator; import org.apache.pinot.controller.helix.core.minion.generator.TaskGeneratorUtils; import org.apache.pinot.core.common.MinionConstants; @@ -96,7 +98,6 @@ public List generateTasks(List tableConfigs) { List notpurgedSegmentsZKMetadata = new ArrayList<>(); for (SegmentZKMetadata segmentMetadata : segmentsZKMetadata) { - if (segmentMetadata.getCustomMap() != null && segmentMetadata.getCustomMap() .containsKey(MinionConstants.PurgeTask.TASK_TYPE + MinionConstants.TASK_TIME_SUFFIX)) { purgedSegmentsZKMetadata.add(segmentMetadata); @@ -113,8 +114,24 @@ public List generateTasks(List tableConfigs) { int tableNumTasks = 0; Set runningSegments = TaskGeneratorUtils.getRunningSegments(MinionConstants.PurgeTask.TASK_TYPE, _clusterInfoAccessor); + List segmentsForDeletion = new ArrayList<>(); + // For realtime tables, build a map of partition to latest segment to avoid deleting last segments + Set lastLLCSegmentPerPartition = new HashSet<>(); + if (tableConfig.getTableType() == TableType.REALTIME) { + lastLLCSegmentPerPartition = getLastLLCSegmentPerPartition(segmentsZKMetadata); + } for (SegmentZKMetadata segmentZKMetadata : notpurgedSegmentsZKMetadata) { String segmentName = segmentZKMetadata.getSegmentName(); + if (segmentZKMetadata.getTotalDocs() == 0L) { + // Check if this empty segment is the last segment of a partition + if (lastLLCSegmentPerPartition.contains(segmentName)) { + LOGGER.info("Skipping deletion of empty segment {} as it is the last segment of its partition", + segmentName); + } else { + segmentsForDeletion.add(segmentName); + } + continue; + } Map configs = new HashMap<>(getBaseTaskConfigs(tableConfig, List.of(segmentName))); Long tsLastPurge; if (segmentZKMetadata.getCustomMap() != null) { @@ -141,9 +158,38 @@ public List generateTasks(List tableConfigs) { pinotTaskConfigs.add(new PinotTaskConfig(taskType, configs)); tableNumTasks++; } + if (!segmentsForDeletion.isEmpty()) { + _clusterInfoAccessor.getPinotHelixResourceManager().deleteSegments(tableName, segmentsForDeletion, + "0d"); + LOGGER.info( + "Deleted segments containing no records for table: {}, number of segments to be deleted: {}", + tableName, segmentsForDeletion.size()); + } LOGGER.info("Finished generating {} tasks configs for table: {} " + "for task: {}", tableNumTasks, tableName, taskType); } return pinotTaskConfigs; } + + private Set getLastLLCSegmentPerPartition(List segmentsZKMetadata) { + Map latestLLCSegmentNameMap = new HashMap<>(); + for (SegmentZKMetadata segmentZKMetadata : segmentsZKMetadata) { + // Skip UPLOADED segments that don't conform to the LLC segment name + LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentZKMetadata.getSegmentName()); + if (llcSegmentName != null) { + latestLLCSegmentNameMap.compute(llcSegmentName.getPartitionGroupId(), (k, latestLLCSegmentName) -> { + if (latestLLCSegmentName == null + || llcSegmentName.getSequenceNumber() > latestLLCSegmentName.getSequenceNumber()) { + return llcSegmentName; + } else { + return latestLLCSegmentName; + } + }); + } + } + Set lastLLCSegmentPerPartition = new HashSet<>(); + latestLLCSegmentNameMap.forEach((ignored, value) -> + lastLLCSegmentPerPartition.add(value.getSegmentName())); + return lastLLCSegmentPerPartition; + } } From 2188ca2b81dc4e83dbad6cbeef528dbe031f349b Mon Sep 17 00:00:00 2001 From: Yash Mayya Date: Fri, 25 Jul 2025 14:38:35 +0530 Subject: [PATCH 024/167] Fix CAST evaluation with literal-only operands in MSQE (#16421) --- .../scalar/DataTypeConversionFunctions.java | 66 +++++++++---------- .../tests/ErrorCodesIntegrationTest.java | 2 +- .../rel/rules/PinotEvaluateLiteralRule.java | 13 +++- .../queries/LiteralEvaluationPlans.json | 32 ++++++++- .../queries/PhysicalOptimizerPlans.json | 26 ++++---- 5 files changed, 88 insertions(+), 51 deletions(-) diff --git a/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/DataTypeConversionFunctions.java b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/DataTypeConversionFunctions.java index 347031ec864d..c18ffed81294 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/DataTypeConversionFunctions.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/DataTypeConversionFunctions.java @@ -38,43 +38,43 @@ private DataTypeConversionFunctions() { @ScalarFunction public static Object cast(Object value, String targetTypeLiteral) { - try { - Class clazz = value.getClass(); - // TODO: Support cast for MV - Preconditions.checkArgument(!clazz.isArray() | clazz == byte[].class, "%s must not be an array type", clazz); - PinotDataType sourceType = PinotDataType.getSingleValueType(clazz); - String transformed = targetTypeLiteral.toUpperCase(); - PinotDataType targetDataType; - switch (transformed) { - case "BIGINT": - targetDataType = LONG; - break; - case "DECIMAL": - targetDataType = BIG_DECIMAL; - break; - case "INT": - targetDataType = INTEGER; - break; - case "VARBINARY": - targetDataType = BYTES; - break; - case "VARCHAR": - targetDataType = STRING; - break; - default: + Class clazz = value.getClass(); + // TODO: Support cast for MV + Preconditions.checkArgument(!clazz.isArray() | clazz == byte[].class, "%s must not be an array type", clazz); + PinotDataType sourceType = PinotDataType.getSingleValueType(clazz); + String transformed = targetTypeLiteral.toUpperCase(); + PinotDataType targetDataType; + switch (transformed) { + case "BIGINT": + targetDataType = LONG; + break; + case "DECIMAL": + targetDataType = BIG_DECIMAL; + break; + case "INT": + targetDataType = INTEGER; + break; + case "VARBINARY": + targetDataType = BYTES; + break; + case "VARCHAR": + targetDataType = STRING; + break; + default: + try { targetDataType = PinotDataType.valueOf(transformed); - break; - } - if (sourceType == STRING && (targetDataType == INTEGER || targetDataType == LONG)) { - if (String.valueOf(value).contains(".")) { - // convert integers via double to avoid parse errors - return targetDataType.convert(DOUBLE.convert(value, sourceType), DOUBLE); + } catch (IllegalArgumentException e) { + throw new IllegalArgumentException("Unknown data type: " + targetTypeLiteral); } + break; + } + if (sourceType == STRING && (targetDataType == INTEGER || targetDataType == LONG)) { + if (String.valueOf(value).contains(".")) { + // convert integers via double to avoid parse errors + return targetDataType.convert(DOUBLE.convert(value, sourceType), DOUBLE); } - return targetDataType.convert(value, sourceType); - } catch (IllegalArgumentException e) { - throw new IllegalArgumentException("Unknown data type: " + targetTypeLiteral); } + return targetDataType.convert(value, sourceType); } /** diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/ErrorCodesIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/ErrorCodesIntegrationTest.java index 383a71fd1d47..069e3fa382d8 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/ErrorCodesIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/ErrorCodesIntegrationTest.java @@ -138,7 +138,7 @@ public void testInvalidCasting() throws Exception { // ArrTime expects a numeric type testQueryException("SELECT COUNT(*) FROM mytable where ArrTime = 'potato'", - useMultiStageQueryEngine() ? QueryErrorCode.QUERY_EXECUTION : QueryErrorCode.QUERY_VALIDATION); + useMultiStageQueryEngine() ? QueryErrorCode.QUERY_PLANNING : QueryErrorCode.QUERY_VALIDATION); } @Test diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotEvaluateLiteralRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotEvaluateLiteralRule.java index ca634a04c2e8..fea69ab5d211 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotEvaluateLiteralRule.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotEvaluateLiteralRule.java @@ -161,15 +161,14 @@ private static RexNode evaluateLiteralOnlyFunction(RexCall rexCall, RexBuilder r assert operands.stream().allMatch( operand -> operand instanceof RexLiteral || (operand instanceof RexCall && ((RexCall) operand).getOperands() .stream().allMatch(op -> op instanceof RexLiteral))); + int numArguments = operands.size(); ColumnDataType[] argumentTypes = new ColumnDataType[numArguments]; Object[] arguments = new Object[numArguments]; for (int i = 0; i < numArguments; i++) { RexNode rexNode = operands.get(i); RexLiteral rexLiteral; - if (rexNode instanceof RexCall && ((RexCall) rexNode).getOperator().getKind() == SqlKind.CAST) { - rexLiteral = (RexLiteral) ((RexCall) rexNode).getOperands().get(0); - } else if (rexNode instanceof RexLiteral) { + if (rexNode instanceof RexLiteral) { rexLiteral = (RexLiteral) rexNode; } else { // Function operands cannot be evaluated, skip @@ -178,6 +177,14 @@ private static RexNode evaluateLiteralOnlyFunction(RexCall rexCall, RexBuilder r argumentTypes[i] = RelToPlanNodeConverter.convertToColumnDataType(rexLiteral.getType()); arguments[i] = getLiteralValue(rexLiteral); } + + if (rexCall.getKind() == SqlKind.CAST) { + // Handle separately because the CAST operator only has one operand (the value to be cast) and the type to be cast + // to is determined by the operator's return type. Pinot's CAST function implementation requires two arguments: + // the value to be cast and the target type. + argumentTypes = new ColumnDataType[]{argumentTypes[0], ColumnDataType.STRING}; + arguments = new Object[]{arguments[0], RelToPlanNodeConverter.convertToColumnDataType(rexCall.getType()).name()}; + } String canonicalName = FunctionRegistry.canonicalize(PinotRuleUtils.extractFunctionName(rexCall)); FunctionInfo functionInfo = FunctionRegistry.lookupFunctionInfo(canonicalName, argumentTypes); if (functionInfo == null) { diff --git a/pinot-query-planner/src/test/resources/queries/LiteralEvaluationPlans.json b/pinot-query-planner/src/test/resources/queries/LiteralEvaluationPlans.json index f23e0662e077..e63b4aab8a68 100644 --- a/pinot-query-planner/src/test/resources/queries/LiteralEvaluationPlans.json +++ b/pinot-query-planner/src/test/resources/queries/LiteralEvaluationPlans.json @@ -45,7 +45,7 @@ "sql": "EXPLAIN PLAN FOR SELECT timestampDiff(DAY, CAST(ts as TIMESTAMP), CAST(dateTrunc('MONTH', FROMDATETIME('1997-02-01 00:00:00', 'yyyy-MM-dd HH:mm:ss')) as TIMESTAMP)) FROM d", "output": [ "Execution Plan", - "\nLogicalProject(EXPR$0=[TIMESTAMPDIFF(FLAG(DAY), CAST($7):TIMESTAMP(0) NOT NULL, CAST(854755200000:BIGINT):TIMESTAMP(0) NOT NULL)])", + "\nLogicalProject(EXPR$0=[TIMESTAMPDIFF(FLAG(DAY), CAST($7):TIMESTAMP(0) NOT NULL, 1997-02-01 00:00:00)])", "\n PinotLogicalTableScan(table=[[default, d]])", "\n" ] @@ -253,6 +253,36 @@ "\n" ] }, + { + "description": "filter with implicit cast", + "sql": "EXPLAIN PLAN FOR SELECT * FROM a WHERE col3 > '10.5'", + "output": [ + "Execution Plan", + "\nLogicalFilter(condition=[>($2, 10)])", + "\n PinotLogicalTableScan(table=[[default, a]])", + "\n" + ] + }, + { + "description": "filter with explicit cast", + "sql": "EXPLAIN PLAN FOR SELECT * FROM a WHERE col3 > CAST('10.5' AS INT)", + "output": [ + "Execution Plan", + "\nLogicalFilter(condition=[>($2, 10)])", + "\n PinotLogicalTableScan(table=[[default, a]])", + "\n" + ] + }, + { + "description": "filter with nested explicit casts", + "sql": "EXPLAIN PLAN FOR SELECT * FROM a WHERE col3 > CAST(CAST('10.5' AS LONG) AS INT)", + "output": [ + "Execution Plan", + "\nLogicalFilter(condition=[>($2, 10)])", + "\n PinotLogicalTableScan(table=[[default, a]])", + "\n" + ] + }, { "description": "select non-exist literal function", "sql": "EXPLAIN PLAN FOR Select nonExistFun(1,2) FROM a", diff --git a/pinot-query-planner/src/test/resources/queries/PhysicalOptimizerPlans.json b/pinot-query-planner/src/test/resources/queries/PhysicalOptimizerPlans.json index ca7dc9730742..106d0e7f9a91 100644 --- a/pinot-query-planner/src/test/resources/queries/PhysicalOptimizerPlans.json +++ b/pinot-query-planner/src/test/resources/queries/PhysicalOptimizerPlans.json @@ -466,7 +466,7 @@ "queries": [ { "description": "Self semi-joins", - "sql": "SET usePhysicalOptimizer=true; EXPLAIN PLAN FOR SELECT col1, col2 FROM a WHERE col2 IN (SELECT col2 FROM a WHERE col3 = 'foo') AND col2 IN (SELECT col2 FROM a WHERE col3 = 'bar')", + "sql": "SET usePhysicalOptimizer=true; EXPLAIN PLAN FOR SELECT col1, col2 FROM a WHERE col2 IN (SELECT col2 FROM a WHERE col3 = '1') AND col2 IN (SELECT col2 FROM a WHERE col3 = '2')", "output": [ "Execution Plan", "\nPhysicalExchange(exchangeStrategy=[SINGLETON_EXCHANGE])", @@ -477,18 +477,18 @@ "\n PhysicalTableScan(table=[[default, a]])", "\n PhysicalExchange(exchangeStrategy=[IDENTITY_EXCHANGE])", "\n PhysicalProject(col2=[$1])", - "\n PhysicalFilter(condition=[=($2, CAST(_UTF-8'foo'):INTEGER NOT NULL)])", + "\n PhysicalFilter(condition=[=($2, 1)])", "\n PhysicalTableScan(table=[[default, a]])", "\n PhysicalExchange(exchangeStrategy=[IDENTITY_EXCHANGE])", "\n PhysicalProject(col2=[$1])", - "\n PhysicalFilter(condition=[=($2, CAST(_UTF-8'bar'):INTEGER NOT NULL)])", + "\n PhysicalFilter(condition=[=($2, 2)])", "\n PhysicalTableScan(table=[[default, a]])", "\n" ] }, { "description": "Self semi and anti semi-joins", - "sql": "SET usePhysicalOptimizer=true; EXPLAIN PLAN FOR SELECT col1, col2 FROM a WHERE col2 IN (SELECT col2 FROM a WHERE col3 = 'foo') AND col2 NOT IN (SELECT col2 FROM a WHERE col3 = 'bar') AND col2 IN (SELECT col2 FROM a WHERE col3 = 'lorem')", + "sql": "SET usePhysicalOptimizer=true; EXPLAIN PLAN FOR SELECT col1, col2 FROM a WHERE col2 IN (SELECT col2 FROM a WHERE col3 = '1') AND col2 NOT IN (SELECT col2 FROM a WHERE col3 = '2') AND col2 IN (SELECT col2 FROM a WHERE col3 = '3')", "output": [ "Execution Plan", "\nPhysicalExchange(exchangeStrategy=[SINGLETON_EXCHANGE])", @@ -503,23 +503,23 @@ "\n PhysicalTableScan(table=[[default, a]])", "\n PhysicalExchange(exchangeStrategy=[IDENTITY_EXCHANGE])", "\n PhysicalProject(col2=[$1])", - "\n PhysicalFilter(condition=[=($2, CAST(_UTF-8'foo'):INTEGER NOT NULL)])", + "\n PhysicalFilter(condition=[=($2, 1)])", "\n PhysicalTableScan(table=[[default, a]])", "\n PhysicalExchange(exchangeStrategy=[IDENTITY_EXCHANGE])", "\n PhysicalAggregate(group=[{0}], agg#0=[MIN($1)], aggType=[DIRECT])", "\n PhysicalProject(col2=[$1], $f1=[true])", - "\n PhysicalFilter(condition=[=($2, CAST(_UTF-8'bar'):INTEGER NOT NULL)])", + "\n PhysicalFilter(condition=[=($2, 2)])", "\n PhysicalTableScan(table=[[default, a]])", "\n PhysicalExchange(exchangeStrategy=[IDENTITY_EXCHANGE])", "\n PhysicalProject(col2=[$1])", - "\n PhysicalFilter(condition=[=($2, CAST(_UTF-8'lorem'):INTEGER NOT NULL)])", + "\n PhysicalFilter(condition=[=($2, 3)])", "\n PhysicalTableScan(table=[[default, a]])", "\n" ] }, { "description": "Self semi and anti semi-joins with aggregation in the end", - "sql": "SET usePhysicalOptimizer=true; EXPLAIN PLAN FOR SELECT col1, COUNT(*) FROM a WHERE col2 IN (SELECT col2 FROM a WHERE col3 = 'foo') AND col2 NOT IN (SELECT col2 FROM a WHERE col3 = 'bar') AND col2 IN (SELECT col2 FROM a WHERE col3 = 'lorem') GROUP BY col1", + "sql": "SET usePhysicalOptimizer=true; EXPLAIN PLAN FOR SELECT col1, COUNT(*) FROM a WHERE col2 IN (SELECT col2 FROM a WHERE col3 = '1') AND col2 NOT IN (SELECT col2 FROM a WHERE col3 = '2') AND col2 IN (SELECT col2 FROM a WHERE col3 = '3') GROUP BY col1", "output": [ "Execution Plan", "\nPhysicalExchange(exchangeStrategy=[SINGLETON_EXCHANGE])", @@ -537,16 +537,16 @@ "\n PhysicalTableScan(table=[[default, a]])", "\n PhysicalExchange(exchangeStrategy=[IDENTITY_EXCHANGE])", "\n PhysicalProject(col2=[$1])", - "\n PhysicalFilter(condition=[=($2, CAST(_UTF-8'foo'):INTEGER NOT NULL)])", + "\n PhysicalFilter(condition=[=($2, 1)])", "\n PhysicalTableScan(table=[[default, a]])", "\n PhysicalExchange(exchangeStrategy=[IDENTITY_EXCHANGE])", "\n PhysicalAggregate(group=[{0}], agg#0=[MIN($1)], aggType=[DIRECT])", "\n PhysicalProject(col2=[$1], $f1=[true])", - "\n PhysicalFilter(condition=[=($2, CAST(_UTF-8'bar'):INTEGER NOT NULL)])", + "\n PhysicalFilter(condition=[=($2, 2)])", "\n PhysicalTableScan(table=[[default, a]])", "\n PhysicalExchange(exchangeStrategy=[IDENTITY_EXCHANGE])", "\n PhysicalProject(col2=[$1])", - "\n PhysicalFilter(condition=[=($2, CAST(_UTF-8'lorem'):INTEGER NOT NULL)])", + "\n PhysicalFilter(condition=[=($2, 3)])", "\n PhysicalTableScan(table=[[default, a]])", "\n" ] @@ -613,7 +613,7 @@ "queries": [ { "description": "Union, distinct, etc. but still maximally identity exchange", - "sql": "SET usePhysicalOptimizer=true; EXPLAIN PLAN FOR WITH tmp AS (SELECT col2 FROM a WHERE col1 = 'foo' UNION ALL SELECT col2 FROM a WHERE col3 = 'bar'), tmp2 AS (SELECT DISTINCT col2 FROM tmp) SELECT COUNT(*), col3 FROM a WHERE col2 IN (SELECT col2 FROM tmp2) GROUP BY col3", + "sql": "SET usePhysicalOptimizer=true; EXPLAIN PLAN FOR WITH tmp AS (SELECT col2 FROM a WHERE col1 = 'foo' UNION ALL SELECT col2 FROM a WHERE col3 = '1'), tmp2 AS (SELECT DISTINCT col2 FROM tmp) SELECT COUNT(*), col3 FROM a WHERE col2 IN (SELECT col2 FROM tmp2) GROUP BY col3", "output": [ "Execution Plan", "\nPhysicalExchange(exchangeStrategy=[SINGLETON_EXCHANGE])", @@ -633,7 +633,7 @@ "\n PhysicalTableScan(table=[[default, a]])", "\n PhysicalExchange(exchangeStrategy=[IDENTITY_EXCHANGE])", "\n PhysicalProject(col2=[$1])", - "\n PhysicalFilter(condition=[=($2, CAST(_UTF-8'bar'):INTEGER NOT NULL)])", + "\n PhysicalFilter(condition=[=($2, 1)])", "\n PhysicalTableScan(table=[[default, a]])", "\n" ] From f7b9a25984467cd884cd9a52c0ab9f79d4cba5e5 Mon Sep 17 00:00:00 2001 From: Praveen Date: Fri, 25 Jul 2025 16:04:43 -0700 Subject: [PATCH 025/167] [Query Resource Isolation] Additonal Sampling for Broker and Server (#16164) * fix * sampling * Broker sampling * revert integ-test * Fix test failures * review comments * remove MSE * broker auth * remove per pruner & planner sample --- .../BaseSingleStageBrokerRequestHandler.java | 7 +++++++ .../core/query/executor/ServerQueryExecutorV1Impl.java | 4 ++++ 2 files changed, 11 insertions(+) diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java index d9a1e365ead6..78b8d4b40405 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java @@ -400,6 +400,8 @@ protected BrokerResponse doHandleRequest(long requestId, String query, SqlNodeAn // full request compile time = compilationTimeNs + parserTimeNs _brokerMetrics.addPhaseTiming(rawTableName, BrokerQueryPhase.REQUEST_COMPILATION, (compilationEndTimeNs - compilationStartTimeNs) + sqlNodeAndOptions.getParseTimeNs()); + // Accounts for resource usage of the compilation phase, since compilation for some queries can be expensive. + Tracing.ThreadAccountantOps.sampleAndCheckInterruption(); // Second-stage table-level access control // TODO: Modify AccessControl interface to directly take PinotQuery @@ -439,6 +441,8 @@ protected BrokerResponse doHandleRequest(long requestId, String query, SqlNodeAn _brokerMetrics.addPhaseTiming(rawTableName, BrokerQueryPhase.AUTHORIZATION, System.nanoTime() - compilationEndTimeNs); + // Accounts for resource usage of the authorization phase. + Tracing.ThreadAccountantOps.sampleAndCheckInterruption(); if (!authorizationResult.hasAccess()) { throwAccessDeniedError(requestId, query, requestContext, tableName, authorizationResult); @@ -687,6 +691,9 @@ protected BrokerResponse doHandleRequest(long requestId, String query, SqlNodeAn long routingEndTimeNs = System.nanoTime(); _brokerMetrics.addPhaseTiming(rawTableName, BrokerQueryPhase.QUERY_ROUTING, routingEndTimeNs - routingStartTimeNs); + // Account the resource used for routing phase, since for single stage queries with multiple segments, routing + // can be expensive. + Tracing.ThreadAccountantOps.sampleAndCheckInterruption(); // Set timeout in the requests long timeSpentMs = TimeUnit.NANOSECONDS.toMillis(routingEndTimeNs - compilationStartTimeNs); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/executor/ServerQueryExecutorV1Impl.java b/pinot-core/src/main/java/org/apache/pinot/core/query/executor/ServerQueryExecutorV1Impl.java index 97ee75dedb46..d63710dd3ffd 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/executor/ServerQueryExecutorV1Impl.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/executor/ServerQueryExecutorV1Impl.java @@ -313,6 +313,8 @@ private InstanceResponseBlock executeInternal(TableExecutionInfo executionInfo, TableExecutionInfo.SelectedSegmentsInfo selectedSegmentsInfo = executionInfo.getSelectedSegmentsInfo(queryContext, timerContext, executorService, _segmentPrunerService); + // Account for resource usage in pruning, given that it can be expensive for large segment lists. + Tracing.ThreadAccountantOps.sampleAndCheckInterruption(); InstanceResponseBlock instanceResponse = execute(selectedSegmentsInfo.getIndexSegments(), queryContext, timerContext, executorService, streamer, @@ -580,6 +582,8 @@ private InstanceResponseBlock executeQuery(QueryContext queryContext, TimerConte } InstanceResponseBlock instanceResponse; Plan queryPlan = planCombineQuery(queryContext, timerContext, executorService, streamer, selectedSegmentContexts); + // Sample to track usage of query planning, since it can be expensive for large segment lists. + Tracing.ThreadAccountantOps.sampleAndCheckInterruption(); TimerContext.Timer planExecTimer = timerContext.startNewPhaseTimer(ServerQueryPhase.QUERY_PLAN_EXECUTION); instanceResponse = queryPlan.execute(); From ee474defe9b1187ff74137847ffd224af65971bf Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Sun, 27 Jul 2025 17:37:07 -0700 Subject: [PATCH 026/167] Bump org.apache.commons:commons-text from 1.13.1 to 1.14.0 (#16427) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index ff94792856c9..5125e92af7b7 100644 --- a/pom.xml +++ b/pom.xml @@ -204,7 +204,7 @@ 3.18.0 4.5.0 - 1.13.1 + 1.14.0 1.27.1 3.6.1 1.14.0 From 313965a628bb626aa6f435bb3b1da16606792176 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Sun, 27 Jul 2025 17:37:28 -0700 Subject: [PATCH 027/167] Bump com.google.errorprone:error_prone_annotations from 2.40.0 to 2.41.0 (#16428) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 5125e92af7b7..2a572cd92375 100644 --- a/pom.xml +++ b/pom.xml @@ -243,7 +243,7 @@ 26.64.0 1.1.1 1.8 - 2.40.0 + 2.41.0 3.0.0 3.0.2 From 7fa52b8b4e129335f7f664495503d5c637353cfb Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Sun, 27 Jul 2025 17:37:45 -0700 Subject: [PATCH 028/167] Bump software.amazon.awssdk:bom from 2.32.7 to 2.32.8 (#16429) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 2a572cd92375..3604138c3353 100644 --- a/pom.xml +++ b/pom.xml @@ -178,7 +178,7 @@ 0.15.1 0.4.7 4.2.2 - 2.32.7 + 2.32.8 1.2.36 1.22.0 2.14.0 From a2557515d92a91257ae165504ffdfabeb981da06 Mon Sep 17 00:00:00 2001 From: Rajat Venkatesh <1638298+vrajat@users.noreply.github.com> Date: Mon, 28 Jul 2025 23:10:58 +0530 Subject: [PATCH 029/167] Use Broker's accountant to sample in the request handler. (#16439) --- .../BaseSingleStageBrokerRequestHandler.java | 8 ++++---- .../src/main/java/org/apache/pinot/spi/trace/Tracing.java | 7 +++++++ 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java index 78b8d4b40405..8ef923f63bbe 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java @@ -332,7 +332,7 @@ protected BrokerResponse handleRequest(long requestId, String query, SqlNodeAndO return doHandleRequest(requestId, query, sqlNodeAndOptions, request, requesterIdentity, requestContext, httpHeaders, accessControl); } finally { - Tracing.ThreadAccountantOps.clear(); + _resourceUsageAccountant.clear(); } } @@ -401,7 +401,7 @@ protected BrokerResponse doHandleRequest(long requestId, String query, SqlNodeAn _brokerMetrics.addPhaseTiming(rawTableName, BrokerQueryPhase.REQUEST_COMPILATION, (compilationEndTimeNs - compilationStartTimeNs) + sqlNodeAndOptions.getParseTimeNs()); // Accounts for resource usage of the compilation phase, since compilation for some queries can be expensive. - Tracing.ThreadAccountantOps.sampleAndCheckInterruption(); + Tracing.ThreadAccountantOps.sampleAndCheckInterruption(_resourceUsageAccountant); // Second-stage table-level access control // TODO: Modify AccessControl interface to directly take PinotQuery @@ -442,7 +442,7 @@ protected BrokerResponse doHandleRequest(long requestId, String query, SqlNodeAn _brokerMetrics.addPhaseTiming(rawTableName, BrokerQueryPhase.AUTHORIZATION, System.nanoTime() - compilationEndTimeNs); // Accounts for resource usage of the authorization phase. - Tracing.ThreadAccountantOps.sampleAndCheckInterruption(); + Tracing.ThreadAccountantOps.sampleAndCheckInterruption(_resourceUsageAccountant); if (!authorizationResult.hasAccess()) { throwAccessDeniedError(requestId, query, requestContext, tableName, authorizationResult); @@ -693,7 +693,7 @@ protected BrokerResponse doHandleRequest(long requestId, String query, SqlNodeAn routingEndTimeNs - routingStartTimeNs); // Account the resource used for routing phase, since for single stage queries with multiple segments, routing // can be expensive. - Tracing.ThreadAccountantOps.sampleAndCheckInterruption(); + Tracing.ThreadAccountantOps.sampleAndCheckInterruption(_resourceUsageAccountant); // Set timeout in the requests long timeSpentMs = TimeUnit.NANOSECONDS.toMillis(routingEndTimeNs - compilationStartTimeNs); diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/trace/Tracing.java b/pinot-spi/src/main/java/org/apache/pinot/spi/trace/Tracing.java index 831e6d08b6b7..0f2ee9deb725 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/trace/Tracing.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/trace/Tracing.java @@ -391,6 +391,13 @@ public static void sampleAndCheckInterruption() { sample(); } + public static void sampleAndCheckInterruption(ThreadResourceUsageAccountant accountant) { + if (Thread.interrupted() || accountant.isAnchorThreadInterrupted() || accountant.isQueryTerminated()) { + throw new EarlyTerminationException("Interrupted while merging records"); + } + accountant.sampleUsage(); + } + @Deprecated public static void updateQueryUsageConcurrently(String queryId) { } From 0bd90a1d4487e307f68eb2241f9e9ce717570a1c Mon Sep 17 00:00:00 2001 From: "Xiaotian (Jackie) Jiang" <17555551+Jackie-Jiang@users.noreply.github.com> Date: Mon, 28 Jul 2025 13:02:09 -0600 Subject: [PATCH 030/167] Enhance and unify JSON index reader (#16436) --- .../impl/json/MutableJsonIndexImpl.java | 405 ++++++++---------- .../readers/BitmapInvertedIndexReader.java | 1 - .../json/ImmutableJsonIndexReader.java | 345 +++++++-------- 3 files changed, 323 insertions(+), 428 deletions(-) diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/json/MutableJsonIndexImpl.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/json/MutableJsonIndexImpl.java index a34d9285d04b..fe6ade159c39 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/json/MutableJsonIndexImpl.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/json/MutableJsonIndexImpl.java @@ -71,18 +71,18 @@ public class MutableJsonIndexImpl implements MutableJsonIndex { private static final Logger LOGGER = LoggerFactory.getLogger(MutableJsonIndexImpl.class); private final JsonIndexConfig _jsonIndexConfig; + private final String _segmentName; + private final String _columnName; private final TreeMap _postingListMap; private final IntList _docIdMapping; + private final long _maxBytesSize; private final ReentrantReadWriteLock.ReadLock _readLock; private final ReentrantReadWriteLock.WriteLock _writeLock; - private final String _segmentName; - private final String _columnName; - private final long _maxBytesSize; + private final ServerMetrics _serverMetrics; private int _nextDocId; private int _nextFlattenedDocId; private long _bytesSize; - private ServerMetrics _serverMetrics; public MutableJsonIndexImpl(JsonIndexConfig jsonIndexConfig, String segmentName, String columnName) { _jsonIndexConfig = jsonIndexConfig; @@ -164,14 +164,17 @@ public MutableRoaringBitmap getMatchingDocIds(Object filterObj) { if (!(filterObj instanceof FilterContext)) { throw new BadQueryRequestException("Invalid json match filter: " + filterObj); } - FilterContext filter = (FilterContext) filterObj; + return getMatchingDocIds((FilterContext) filterObj); + } + private MutableRoaringBitmap getMatchingDocIds(FilterContext filter) { _readLock.lock(); try { - if (filter.getType() == FilterContext.Type.PREDICATE && isExclusive(filter.getPredicate().getType())) { + Predicate predicate = filter.getPredicate(); + if (predicate != null && isExclusive(predicate.getType())) { // Handle exclusive predicate separately because the flip can only be applied to the unflattened doc ids in // order to get the correct result, and it cannot be nested - LazyBitmap flattenedDocIds = getMatchingFlattenedDocIds(filter.getPredicate()); + LazyBitmap flattenedDocIds = getMatchingFlattenedDocIds(predicate); MutableRoaringBitmap matchingDocIds = new MutableRoaringBitmap(); flattenedDocIds.forEach(flattenedDocId -> matchingDocIds.add(_docIdMapping.getInt(flattenedDocId))); matchingDocIds.flip(0, (long) _nextDocId); @@ -198,86 +201,131 @@ private boolean isExclusive(Predicate.Type predicateType) { * It stores either a bitmap from posting list that must be cloned before mutating (readOnly=true) * or an already cloned bitmap. */ - static class LazyBitmap { - - final static LazyBitmap EMPTY_BITMAP = new LazyBitmap(null); + private static class LazyBitmap { + final static LazyBitmap EMPTY_BITMAP = createImmutable(new RoaringBitmap()); // value should be null only for EMPTY - @Nullable - RoaringBitmap _value; + final RoaringBitmap _value; // if readOnly then bitmap needs to be cloned before applying mutating operations - boolean _readOnly; + final boolean _mutable; - LazyBitmap(RoaringBitmap bitmap) { + LazyBitmap(RoaringBitmap bitmap, boolean mutable) { _value = bitmap; - _readOnly = true; + _mutable = mutable; } - LazyBitmap(RoaringBitmap bitmap, boolean isReadOnly) { - _value = bitmap; - _readOnly = isReadOnly; + static LazyBitmap createMutable(RoaringBitmap bitmap) { + return new LazyBitmap(bitmap, true); + } + + static LazyBitmap createImmutable(RoaringBitmap bitmap) { + return new LazyBitmap(bitmap, false); } boolean isMutable() { - return !_readOnly; + return _mutable; } LazyBitmap toMutable() { - if (_readOnly) { - if (_value == null) { - return new LazyBitmap(new RoaringBitmap(), false); - } - - _value = _value.clone(); - _readOnly = false; + if (_mutable) { + return this; } - - return this; + return createMutable(_value.clone()); } - void and(LazyBitmap bitmap) { - assert isMutable(); - - _value.and(bitmap._value); + LazyBitmap and(LazyBitmap other) { + if (isEmpty() || other.isEmpty()) { + return LazyBitmap.EMPTY_BITMAP; + } + if (isMutable()) { + _value.and(other._value); + return this; + } + if (other.isMutable()) { + other._value.and(_value); + return other; + } + return createMutable(RoaringBitmap.and(_value, other._value)); } LazyBitmap and(RoaringBitmap bitmap) { - LazyBitmap mutable = toMutable(); - mutable._value.and(bitmap); - return mutable; + if (isEmpty() || bitmap.isEmpty()) { + return EMPTY_BITMAP; + } + if (isMutable()) { + _value.and(bitmap); + return this; + } + return createMutable(RoaringBitmap.and(_value, bitmap)); } - LazyBitmap andNot(RoaringBitmap bitmap) { - LazyBitmap mutable = toMutable(); - mutable._value.andNot(bitmap); - return mutable; + LazyBitmap or(LazyBitmap other) { + if (isEmpty()) { + return other; + } + if (other.isEmpty()) { + return this; + } + if (isMutable()) { + _value.or(other._value); + return this; + } + if (other.isMutable()) { + other._value.or(_value); + return other; + } + return createMutable(RoaringBitmap.or(_value, other._value)); } - void or(LazyBitmap bitmap) { - assert isMutable(); + LazyBitmap or(RoaringBitmap bitmap) { + if (isEmpty()) { + return createImmutable(bitmap); + } + if (bitmap.isEmpty()) { + return this; + } + if (isMutable()) { + _value.or(bitmap); + return this; + } + return createMutable(RoaringBitmap.or(_value, bitmap)); + } - _value.or(bitmap._value); + LazyBitmap andNot(LazyBitmap other) { + if (isEmpty()) { + return EMPTY_BITMAP; + } + if (other.isEmpty()) { + return this; + } + if (isMutable()) { + _value.andNot(other._value); + return this; + } + return createMutable(RoaringBitmap.andNot(_value, other._value)); } - LazyBitmap or(RoaringBitmap bitmap) { - LazyBitmap mutable = toMutable(); - mutable._value.or(bitmap); - return mutable; + LazyBitmap andNot(RoaringBitmap bitmap) { + if (isEmpty()) { + return EMPTY_BITMAP; + } + if (bitmap.isEmpty()) { + return this; + } + if (isMutable()) { + _value.andNot(bitmap); + return this; + } + return createMutable(RoaringBitmap.andNot(_value, bitmap)); } boolean isEmpty() { - if (_value == null) { - return true; - } else { - return _value.isEmpty(); - } + return _value.isEmpty(); } void forEach(IntConsumer ic) { - if (_value != null) { - _value.forEach(ic); - } + _value.forEach(ic); } LazyBitmap flip(long rangeStart, long rangeEnd) { @@ -287,11 +335,7 @@ LazyBitmap flip(long rangeStart, long rangeEnd) { } RoaringBitmap getValue() { - if (_value == null) { - return new RoaringBitmap(); - } else { - return _value; - } + return _value; } } @@ -305,30 +349,19 @@ private LazyBitmap getMatchingFlattenedDocIds(FilterContext filter) { LazyBitmap matchingDocIds = getMatchingFlattenedDocIds(filters.get(0)); for (int i = 1, numFilters = filters.size(); i < numFilters; i++) { if (matchingDocIds.isEmpty()) { - break; + return LazyBitmap.EMPTY_BITMAP; } - LazyBitmap filterDocIds = getMatchingFlattenedDocIds(filters.get(i)); - if (filterDocIds.isEmpty()) { - return filterDocIds; - } else { - matchingDocIds = and(matchingDocIds, filterDocIds); - } + matchingDocIds = matchingDocIds.and(filterDocIds); } return matchingDocIds; } case OR: { List filters = filter.getChildren(); LazyBitmap matchingDocIds = getMatchingFlattenedDocIds(filters.get(0)); - for (int i = 1, numFilters = filters.size(); i < numFilters; i++) { LazyBitmap filterDocIds = getMatchingFlattenedDocIds(filters.get(i)); - // avoid having to convert matchingDocIds to mutable map - if (filterDocIds.isEmpty()) { - continue; - } - - matchingDocIds = or(matchingDocIds, filterDocIds); + matchingDocIds = matchingDocIds.or(filterDocIds); } return matchingDocIds; } @@ -353,109 +386,100 @@ private LazyBitmap getMatchingFlattenedDocIds(Predicate predicate) { Preconditions.checkArgument(lhs.getType() == ExpressionContext.Type.IDENTIFIER, "Left-hand side of the predicate must be an identifier, got: %s (%s). Put double quotes around the identifier" + " if needed.", lhs, lhs.getType()); - String key = lhs.getIdentifier(); // Support 2 formats: // - JSONPath format (e.g. "$.a[1].b"='abc', "$[0]"=1, "$"='abc') // - Legacy format (e.g. "a[1].b"='abc') + String key = lhs.getIdentifier(); if (key.charAt(0) == '$') { key = key.substring(1); } else { key = JsonUtils.KEY_SEPARATOR + key; } + Pair pair = getKeyAndFlattenedDocIds(key); key = pair.getLeft(); - LazyBitmap matchingDocIds = pair.getRight(); - if (matchingDocIds != null && matchingDocIds.isEmpty()) { + LazyBitmap matchingDocIdsForKey = pair.getRight(); + if (matchingDocIdsForKey != null && matchingDocIdsForKey.isEmpty()) { return LazyBitmap.EMPTY_BITMAP; } + LazyBitmap matchingDocIdsForKeyValue = getMatchingFlattenedDocIdsForKeyValue(predicate, key); + if (matchingDocIdsForKey == null) { + return matchingDocIdsForKeyValue; + } else { + return matchingDocIdsForKeyValue.and(matchingDocIdsForKey); + } + } + private LazyBitmap getMatchingFlattenedDocIdsForKeyValue(Predicate predicate, String key) { Predicate.Type predicateType = predicate.getType(); switch (predicateType) { case EQ: { String value = ((EqPredicate) predicate).getValue(); - String keyValuePair = key + JsonIndexCreator.KEY_VALUE_SEPARATOR + value; - RoaringBitmap result = _postingListMap.get(keyValuePair); - return filter(result, matchingDocIds); + RoaringBitmap docIds = _postingListMap.get(key + JsonIndexCreator.KEY_VALUE_SEPARATOR + value); + return docIds != null ? LazyBitmap.createImmutable(docIds) : LazyBitmap.EMPTY_BITMAP; } case NOT_EQ: { - String notEqualValue = ((NotEqPredicate) predicate).getValue(); - LazyBitmap result = null; - RoaringBitmap allDocIds = _postingListMap.get(key); - if (allDocIds != null && !allDocIds.isEmpty()) { - result = new LazyBitmap(allDocIds); - - RoaringBitmap notEqualDocIds = - _postingListMap.get(key + JsonIndexCreator.KEY_VALUE_SEPARATOR + notEqualValue); - - if (notEqualDocIds != null && !notEqualDocIds.isEmpty()) { - result = result.andNot(notEqualDocIds); - } + if (allDocIds == null) { + return LazyBitmap.EMPTY_BITMAP; + } + LazyBitmap result = LazyBitmap.createImmutable(allDocIds); + String value = ((NotEqPredicate) predicate).getValue(); + RoaringBitmap docIds = _postingListMap.get(key + JsonIndexCreator.KEY_VALUE_SEPARATOR + value); + if (docIds != null) { + return result.andNot(docIds); + } else { + return result; } - - return filter(result, matchingDocIds); } case IN: { - List values = ((InPredicate) predicate).getValues(); - LazyBitmap result = null; - StringBuilder buffer = new StringBuilder(key); buffer.append(JsonIndexCreator.KEY_VALUE_SEPARATOR); int pos = buffer.length(); - + LazyBitmap result = LazyBitmap.EMPTY_BITMAP; + List values = ((InPredicate) predicate).getValues(); for (String value : values) { buffer.setLength(pos); buffer.append(value); - String keyValue = buffer.toString(); - - RoaringBitmap docIds = _postingListMap.get(keyValue); - - if (docIds != null && !docIds.isEmpty()) { - if (result == null) { - result = new LazyBitmap(docIds); - } else { - result = result.or(docIds); - } + RoaringBitmap docIds = _postingListMap.get(buffer.toString()); + if (docIds != null) { + result = result.or(docIds); } } - - return filter(result, matchingDocIds); + return result; } case NOT_IN: { - List notInValues = ((NotInPredicate) predicate).getValues(); - LazyBitmap result = null; - RoaringBitmap allDocIds = _postingListMap.get(key); - if (allDocIds != null && !allDocIds.isEmpty()) { - result = new LazyBitmap(allDocIds); - - StringBuilder buffer = new StringBuilder(key); - buffer.append(JsonIndexCreator.KEY_VALUE_SEPARATOR); - int pos = buffer.length(); - - for (String notInValue : notInValues) { - buffer.setLength(pos); - buffer.append(notInValue); - String keyValuePair = buffer.toString(); - - RoaringBitmap docIds = _postingListMap.get(keyValuePair); - if (docIds != null && !docIds.isEmpty()) { - result = result.andNot(docIds); - } + if (allDocIds == null) { + return LazyBitmap.EMPTY_BITMAP; + } + StringBuilder buffer = new StringBuilder(key); + buffer.append(JsonIndexCreator.KEY_VALUE_SEPARATOR); + int pos = buffer.length(); + LazyBitmap result = LazyBitmap.createImmutable(allDocIds); + List values = ((NotInPredicate) predicate).getValues(); + for (String value : values) { + if (result.isEmpty()) { + return LazyBitmap.EMPTY_BITMAP; + } + buffer.setLength(pos); + buffer.append(value); + RoaringBitmap docIds = _postingListMap.get(buffer.toString()); + if (docIds != null) { + result = result.andNot(docIds); } } - - return filter(result, matchingDocIds); + return result; } case IS_NOT_NULL: case IS_NULL: { - RoaringBitmap result = _postingListMap.get(key); - return filter(result, matchingDocIds); + RoaringBitmap docIds = _postingListMap.get(key); + return docIds != null ? LazyBitmap.createImmutable(docIds) : LazyBitmap.EMPTY_BITMAP; } case REGEXP_LIKE: { @@ -463,32 +487,20 @@ private LazyBitmap getMatchingFlattenedDocIds(Predicate predicate) { if (subMap.isEmpty()) { return LazyBitmap.EMPTY_BITMAP; } - Pattern pattern = ((RegexpLikePredicate) predicate).getPattern(); Matcher matcher = pattern.matcher(""); - LazyBitmap result = null; + LazyBitmap result = LazyBitmap.EMPTY_BITMAP; StringBuilder value = new StringBuilder(); - + int valueStart = key.length() + 1; for (Map.Entry entry : subMap.entrySet()) { String keyValue = entry.getKey(); value.setLength(0); - value.append(keyValue, key.length() + 1, keyValue.length()); - - if (!matcher.reset(value).matches()) { - continue; - } - - RoaringBitmap docIds = entry.getValue(); - if (docIds != null && !docIds.isEmpty()) { - if (result == null) { - result = new LazyBitmap(docIds); - } else { - result = result.or(docIds); - } + value.append(keyValue, valueStart, keyValue.length()); + if (matcher.reset(value).matches()) { + result = result.or(entry.getValue()); } } - - return filter(result, matchingDocIds); + return result; } case RANGE: { @@ -496,8 +508,6 @@ private LazyBitmap getMatchingFlattenedDocIds(Predicate predicate) { if (subMap.isEmpty()) { return LazyBitmap.EMPTY_BITMAP; } - - LazyBitmap result = null; RangePredicate rangePredicate = (RangePredicate) predicate; FieldSpec.DataType rangeDataType = rangePredicate.getRangeDataType(); // Simplify to only support numeric and string types @@ -506,16 +516,16 @@ private LazyBitmap getMatchingFlattenedDocIds(Predicate predicate) { } else { rangeDataType = FieldSpec.DataType.STRING; } - boolean lowerUnbounded = rangePredicate.getLowerBound().equals(RangePredicate.UNBOUNDED); boolean upperUnbounded = rangePredicate.getUpperBound().equals(RangePredicate.UNBOUNDED); boolean lowerInclusive = lowerUnbounded || rangePredicate.isLowerInclusive(); boolean upperInclusive = upperUnbounded || rangePredicate.isUpperInclusive(); Object lowerBound = lowerUnbounded ? null : rangeDataType.convert(rangePredicate.getLowerBound()); Object upperBound = upperUnbounded ? null : rangeDataType.convert(rangePredicate.getUpperBound()); - + LazyBitmap result = LazyBitmap.EMPTY_BITMAP; + int valueStart = key.length() + 1; for (Map.Entry entry : subMap.entrySet()) { - Object valueObj = rangeDataType.convert(entry.getKey().substring(key.length() + 1)); + Object valueObj = rangeDataType.convert(entry.getKey().substring(valueStart)); boolean lowerCompareResult = lowerUnbounded || (lowerInclusive ? rangeDataType.compare(valueObj, lowerBound) >= 0 : rangeDataType.compare(valueObj, lowerBound) > 0); @@ -523,15 +533,10 @@ private LazyBitmap getMatchingFlattenedDocIds(Predicate predicate) { upperUnbounded || (upperInclusive ? rangeDataType.compare(valueObj, upperBound) <= 0 : rangeDataType.compare(valueObj, upperBound) < 0); if (lowerCompareResult && upperCompareResult) { - if (result == null) { - result = new LazyBitmap(entry.getValue()); - } else { - result = result.or(entry.getValue()); - } + result = result.or(entry.getValue()); } } - - return filter(result, matchingDocIds); + return result; } default: @@ -651,7 +656,7 @@ private Pair getKeyAndFlattenedDocIds(String key) { if (docIds != null) { if (matchingDocIds == null) { - matchingDocIds = new LazyBitmap(docIds); + matchingDocIds = LazyBitmap.createImmutable(docIds); } else { matchingDocIds = matchingDocIds.and(docIds); } @@ -671,8 +676,7 @@ private Map getMatchingKeysMap(String key) { } @Override - public String[][] getValuesMV(int[] docIds, int length, - Map valueToMatchingFlattenedDocs) { + public String[][] getValuesMV(int[] docIds, int length, Map valueToMatchingFlattenedDocs) { String[][] result = new String[length][]; List>> docIdToFlattenedDocIdsAndValues = new ArrayList<>(); for (int i = 0; i < length; i++) { @@ -751,6 +755,11 @@ public String[] getValuesSV(int[] docIds, int length, Map return values; } + @Override + public boolean canAddMore() { + return _bytesSize < _maxBytesSize; + } + @Override public void close() { try { @@ -763,70 +772,4 @@ public void close() { _segmentName, _columnName, _bytesSize, e); } } - - @Override - public boolean canAddMore() { - return _bytesSize < _maxBytesSize; - } - - // AND given bitmaps, optionally converting first one to mutable (if it's not already) - private static LazyBitmap and(LazyBitmap target, LazyBitmap other) { - if (target.isMutable()) { - target.and(other); - return target; - } else if (other.isMutable()) { - other.and(target); - return other; - } else { - LazyBitmap mutableTarget = target.toMutable(); - mutableTarget.and(other); - return mutableTarget; - } - } - - private static LazyBitmap and(LazyBitmap target, RoaringBitmap other) { - if (target.isMutable()) { - target.and(other); - return target; - } else { - LazyBitmap mutableTarget = target.toMutable(); - mutableTarget.and(other); - return mutableTarget; - } - } - - // OR given bitmaps, optionally converting first one to mutable (if it's not already) - private static LazyBitmap or(LazyBitmap target, LazyBitmap other) { - if (target.isMutable()) { - target.or(other); - return target; - } else if (other.isMutable()) { - other.or(target); - return other; - } else { - LazyBitmap mutableTarget = target.toMutable(); - mutableTarget.or(other); - return mutableTarget; - } - } - - private static LazyBitmap filter(LazyBitmap result, LazyBitmap matchingDocIds) { - if (result == null) { - return LazyBitmap.EMPTY_BITMAP; - } else if (matchingDocIds == null) { - return result; - } else { - return and(matchingDocIds, result); - } - } - - private static LazyBitmap filter(RoaringBitmap result, LazyBitmap matchingDocIds) { - if (result == null) { - return LazyBitmap.EMPTY_BITMAP; - } else if (matchingDocIds == null) { - return new LazyBitmap(result); - } else { - return and(matchingDocIds, result); - } - } } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/BitmapInvertedIndexReader.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/BitmapInvertedIndexReader.java index af678b001864..74a84ca35ecf 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/BitmapInvertedIndexReader.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/BitmapInvertedIndexReader.java @@ -49,7 +49,6 @@ public BitmapInvertedIndexReader(PinotDataBuffer dataBuffer, int numBitmaps) { _firstOffset = getOffset(0); } - @SuppressWarnings("unchecked") @Override public ImmutableRoaringBitmap getDocIds(int dictId) { long offset = getOffset(dictId); diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/json/ImmutableJsonIndexReader.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/json/ImmutableJsonIndexReader.java index 596c426f6f12..20c950d747d1 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/json/ImmutableJsonIndexReader.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/readers/json/ImmutableJsonIndexReader.java @@ -58,7 +58,6 @@ import org.apache.pinot.spi.utils.JsonUtils; import org.apache.pinot.sql.parsers.CalciteSqlParser; import org.roaringbitmap.IntConsumer; -import org.roaringbitmap.PeekableIntIterator; import org.roaringbitmap.RoaringBitmap; import org.roaringbitmap.buffer.ImmutableRoaringBitmap; import org.roaringbitmap.buffer.MutableRoaringBitmap; @@ -136,11 +135,15 @@ public MutableRoaringBitmap getMatchingDocIds(Object filterObj) { if (!(filterObj instanceof FilterContext)) { throw new BadQueryRequestException("Invalid json match filter: " + filterObj); } - FilterContext filter = (FilterContext) filterObj; - if (filter.getType() == FilterContext.Type.PREDICATE && isExclusive(filter.getPredicate().getType())) { + return getMatchingDocIds((FilterContext) filterObj); + } + + private MutableRoaringBitmap getMatchingDocIds(FilterContext filter) { + Predicate predicate = filter.getPredicate(); + if (predicate != null && isExclusive(predicate.getType())) { // Handle exclusive predicate separately because the flip can only be applied to the unflattened doc ids in order // to get the correct result, and it cannot be nested - ImmutableRoaringBitmap flattenedDocIds = getMatchingFlattenedDocIds(filter.getPredicate()); + ImmutableRoaringBitmap flattenedDocIds = getMatchingFlattenedDocIds(predicate); MutableRoaringBitmap resultDocIds = new MutableRoaringBitmap(); flattenedDocIds.forEach((IntConsumer) flattenedDocId -> resultDocIds.add(getDocId(flattenedDocId))); resultDocIds.flip(0, _numDocs); @@ -160,59 +163,51 @@ private boolean isExclusive(Predicate.Type predicateType) { return predicateType == Predicate.Type.IS_NULL; } - // AND given bitmaps, optionally converting first one to mutable (if it's not already) - private static MutableRoaringBitmap and(ImmutableRoaringBitmap target, ImmutableRoaringBitmap other) { + private static ImmutableRoaringBitmap and(ImmutableRoaringBitmap target, ImmutableRoaringBitmap other) { + if (target.isEmpty() || other.isEmpty()) { + return EMPTY_BITMAP; + } if (target instanceof MutableRoaringBitmap) { - MutableRoaringBitmap mutableTarget = (MutableRoaringBitmap) target; - mutableTarget.and(other); - return mutableTarget; - } else if (other instanceof MutableRoaringBitmap) { - MutableRoaringBitmap mutableOther = (MutableRoaringBitmap) other; - mutableOther.and(target); - return mutableOther; - } else { // base implementation - MutableRoaringBitmap mutableTarget = toMutable(target); - mutableTarget.and(other); - return mutableTarget; + ((MutableRoaringBitmap) target).and(other); + return target; } - } - - private static ImmutableRoaringBitmap filter(ImmutableRoaringBitmap result, - ImmutableRoaringBitmap matchingDocIds) { - if (result == null) { - return EMPTY_BITMAP; - } else if (matchingDocIds == null) { - return result; - } else { - return and(matchingDocIds, result); + if (other instanceof MutableRoaringBitmap) { + ((MutableRoaringBitmap) other).and(target); + return other; } + return ImmutableRoaringBitmap.and(target, other); } - // OR given bitmaps, optionally converting first one to mutable (if it's not already) - private static MutableRoaringBitmap or(ImmutableRoaringBitmap target, ImmutableRoaringBitmap other) { + private static ImmutableRoaringBitmap or(ImmutableRoaringBitmap target, ImmutableRoaringBitmap other) { + if (target.isEmpty()) { + return other; + } + if (other.isEmpty()) { + return target; + } if (target instanceof MutableRoaringBitmap) { - MutableRoaringBitmap mutableTarget = (MutableRoaringBitmap) target; - mutableTarget.or(other); - return mutableTarget; - } else if (other instanceof MutableRoaringBitmap) { - MutableRoaringBitmap mutableOther = (MutableRoaringBitmap) other; - mutableOther.or(target); - return mutableOther; - } else { // base implementation - MutableRoaringBitmap mutableTarget = toMutable(target); - mutableTarget.or(other); - return mutableTarget; + ((MutableRoaringBitmap) target).or(other); + return target; + } + if (other instanceof MutableRoaringBitmap) { + ((MutableRoaringBitmap) other).or(target); + return other; } + return ImmutableRoaringBitmap.or(target, other); } - // If given bitmap is not mutable, convert it to such - // used to delay immutable -> mutable conversion as much as possible - private static MutableRoaringBitmap toMutable(ImmutableRoaringBitmap bitmap) { - if (bitmap instanceof MutableRoaringBitmap) { - return (MutableRoaringBitmap) bitmap; - } else { - return bitmap.toMutableRoaringBitmap(); + private static ImmutableRoaringBitmap andNot(ImmutableRoaringBitmap target, ImmutableRoaringBitmap other) { + if (target.isEmpty()) { + return EMPTY_BITMAP; + } + if (other.isEmpty()) { + return target; } + if (target instanceof MutableRoaringBitmap) { + ((MutableRoaringBitmap) target).andNot(other); + return target; + } + return ImmutableRoaringBitmap.andNot(target, other); } /** @@ -223,41 +218,29 @@ private ImmutableRoaringBitmap getMatchingFlattenedDocIds(FilterContext filter) case AND: { List filters = filter.getChildren(); ImmutableRoaringBitmap matchingDocIds = getMatchingFlattenedDocIds(filters.get(0)); - for (int i = 1, numFilters = filters.size(); i < numFilters; i++) { // if current set is empty then there is no point AND-ing it with another one if (matchingDocIds.isEmpty()) { - break; + return EMPTY_BITMAP; } - ImmutableRoaringBitmap filterDocIds = getMatchingFlattenedDocIds(filters.get(i)); - if (filterDocIds.isEmpty()) { - // potentially avoid converting matchingDocIds to mutable map - return filterDocIds; - } else { - matchingDocIds = and(matchingDocIds, filterDocIds); - } + matchingDocIds = and(matchingDocIds, filterDocIds); } return matchingDocIds; } case OR: { List filters = filter.getChildren(); ImmutableRoaringBitmap matchingDocIds = getMatchingFlattenedDocIds(filters.get(0)); - for (int i = 1, numFilters = filters.size(); i < numFilters; i++) { ImmutableRoaringBitmap filterDocIds = getMatchingFlattenedDocIds(filters.get(i)); - // avoid having to convert matchingDocIds to mutable map - if (filterDocIds.isEmpty()) { - continue; - } matchingDocIds = or(matchingDocIds, filterDocIds); } return matchingDocIds; } case PREDICATE: { Predicate predicate = filter.getPredicate(); - Preconditions - .checkArgument(!isExclusive(predicate.getType()), "Exclusive predicate: %s cannot be nested", predicate); + Preconditions.checkArgument(!isExclusive(predicate.getType()), "Exclusive predicate: %s cannot be nested", + predicate); return getMatchingFlattenedDocIds(predicate); } default: @@ -276,10 +259,11 @@ private ImmutableRoaringBitmap getMatchingFlattenedDocIds(Predicate predicate) { Preconditions.checkArgument(lhs.getType() == ExpressionContext.Type.IDENTIFIER, "Left-hand side of the predicate must be an identifier, got: %s (%s). Put double quotes around the identifier" + " if needed.", lhs, lhs.getType()); - String key = lhs.getIdentifier(); + // Support 2 formats: // - JSONPath format (e.g. "$.a[1].b"='abc', "$[0]"=1, "$"='abc') // - Legacy format (e.g. "a[1].b"='abc') + String key = lhs.getIdentifier(); if (_version == BaseJsonIndexCreator.VERSION_2) { if (key.startsWith("$")) { key = key.substring(1); @@ -292,176 +276,160 @@ private ImmutableRoaringBitmap getMatchingFlattenedDocIds(Predicate predicate) { key = key.substring(2); } } + Pair pair = getKeyAndFlattenedDocIds(key); key = pair.getLeft(); - ImmutableRoaringBitmap matchingDocIds = pair.getRight(); - if (matchingDocIds != null && matchingDocIds.isEmpty()) { - return matchingDocIds; + ImmutableRoaringBitmap matchingDocIdsForKey = pair.getRight(); + if (matchingDocIdsForKey != null && matchingDocIdsForKey.isEmpty()) { + return EMPTY_BITMAP; + } + ImmutableRoaringBitmap matchingDocIdsForKeyValue = getMatchingFlattenedDocIdsForKeyValue(predicate, key); + if (matchingDocIdsForKey == null) { + return matchingDocIdsForKeyValue; + } else { + return and(matchingDocIdsForKeyValue, matchingDocIdsForKey); } + } + private ImmutableRoaringBitmap getMatchingFlattenedDocIdsForKeyValue(Predicate predicate, String key) { Predicate.Type predicateType = predicate.getType(); switch (predicateType) { case EQ: { String value = ((EqPredicate) predicate).getValue(); - String keyValuePair = key + JsonIndexCreator.KEY_VALUE_SEPARATOR + value; - int dictId = _dictionary.indexOf(keyValuePair); - ImmutableRoaringBitmap result = null; + int dictId = _dictionary.indexOf(key + JsonIndexCreator.KEY_VALUE_SEPARATOR + value); if (dictId >= 0) { - result = _invertedIndex.getDocIds(dictId); + return _invertedIndex.getDocIds(dictId); + } else { + return EMPTY_BITMAP; } - return filter(result, matchingDocIds); } case NOT_EQ: { - // each array is un-nested and so flattened json document contains only one value - // that means for each key-value pair the set of flattened document ids is disjoint - String notEqualValue = ((NotEqPredicate) predicate).getValue(); - ImmutableRoaringBitmap result = null; - // read bitmap with all values for this key instead of OR-ing many per-value bitmaps int allValuesDictId = _dictionary.indexOf(key); - if (allValuesDictId >= 0) { - ImmutableRoaringBitmap allValuesDocIds = _invertedIndex.getDocIds(allValuesDictId); - - if (!allValuesDocIds.isEmpty()) { - int notEqDictId = _dictionary.indexOf(key + JsonIndexCreator.KEY_VALUE_SEPARATOR + notEqualValue); - if (notEqDictId >= 0) { - ImmutableRoaringBitmap notEqDocIds = _invertedIndex.getDocIds(notEqDictId); - if (notEqDocIds.isEmpty()) { - // there's no value to remove, use found bitmap (is this possible ?) - result = allValuesDocIds; - } else { - // remove doc ids for unwanted value - MutableRoaringBitmap mutableBitmap = allValuesDocIds.toMutableRoaringBitmap(); - mutableBitmap.andNot(notEqDocIds); - result = mutableBitmap; - } - } else { // there's no value to remove, use found bitmap - result = allValuesDocIds; - } - } + if (allValuesDictId < 0) { + return EMPTY_BITMAP; + } + ImmutableRoaringBitmap allValuesDocIds = _invertedIndex.getDocIds(allValuesDictId); + String value = ((NotEqPredicate) predicate).getValue(); + int dictId = _dictionary.indexOf(key + JsonIndexCreator.KEY_VALUE_SEPARATOR + value); + if (dictId >= 0) { + return andNot(allValuesDocIds, _invertedIndex.getDocIds(dictId)); + } else { + // there's no value to remove, use found bitmap + return allValuesDocIds; } - - return filter(result, matchingDocIds); } case IN: { + StringBuilder buffer = new StringBuilder(key); + buffer.append(JsonIndexCreator.KEY_VALUE_SEPARATOR); + int pos = buffer.length(); + ImmutableRoaringBitmap result = EMPTY_BITMAP; List values = ((InPredicate) predicate).getValues(); - ImmutableRoaringBitmap result = null; for (String value : values) { - String keyValuePair = key + JsonIndexCreator.KEY_VALUE_SEPARATOR + value; - int dictId = _dictionary.indexOf(keyValuePair); + buffer.setLength(pos); + buffer.append(value); + int dictId = _dictionary.indexOf(buffer.toString()); if (dictId >= 0) { - ImmutableRoaringBitmap docIds = _invertedIndex.getDocIds(dictId); - if (result == null) { - result = docIds; - } else { - result = or(result, docIds); - } + result = or(result, _invertedIndex.getDocIds(dictId)); } } - - return filter(result, matchingDocIds); + return result; } case NOT_IN: { - List notInValues = ((NotInPredicate) predicate).getValues(); - int[] dictIds = getDictIdRangeForKey(key); - ImmutableRoaringBitmap result = null; - - int valueCount = dictIds[1] - dictIds[0]; - - if (notInValues.size() < valueCount / 2) { + int[] dictIdRange = getDictIdRangeForKey(key); + int minDictId = dictIdRange[0]; + if (minDictId < 0) { + return EMPTY_BITMAP; + } + StringBuilder buffer = new StringBuilder(key); + buffer.append(JsonIndexCreator.KEY_VALUE_SEPARATOR); + int pos = buffer.length(); + int valueCount = dictIdRange[1] - minDictId; + List values = ((NotInPredicate) predicate).getValues(); + if (values.size() < valueCount / 2) { // if there is less notIn values than In values // read bitmap for all values and then remove values from bitmaps associated with notIn values - int allValuesDictId = _dictionary.indexOf(key); - if (allValuesDictId >= 0) { - ImmutableRoaringBitmap allValuesDocIds = _invertedIndex.getDocIds(allValuesDictId); - - if (!allValuesDocIds.isEmpty()) { - result = allValuesDocIds; - - for (String notInValue : notInValues) { - int notInDictId = _dictionary.indexOf(key + JsonIndexCreator.KEY_VALUE_SEPARATOR + notInValue); - if (notInDictId >= 0) { - ImmutableRoaringBitmap notEqDocIds = _invertedIndex.getDocIds(notInDictId); - // remove doc ids for unwanted value - MutableRoaringBitmap mutableBitmap = toMutable(result); - mutableBitmap.andNot(notEqDocIds); - result = mutableBitmap; - } - } + int allValuesDictId = minDictId - 1; + ImmutableRoaringBitmap result = _invertedIndex.getDocIds(allValuesDictId); + for (String value : values) { + if (result.isEmpty()) { + return EMPTY_BITMAP; + } + buffer.setLength(pos); + buffer.append(value); + int dictId = _dictionary.indexOf(buffer.toString()); + if (dictId >= 0) { + // remove doc ids for unwanted value + result = andNot(result, _invertedIndex.getDocIds(dictId)); } } + return result; } else { // if there is more In values than notIn then OR bitmaps for all values except notIn values // resolve dict ids for string values to avoid comparing strings - IntOpenHashSet notInDictIds = null; - if (dictIds[0] < dictIds[1]) { - notInDictIds = new IntOpenHashSet(); - for (String notInValue : notInValues) { - int dictId = _dictionary.indexOf(key + JsonIndexCreator.KEY_VALUE_SEPARATOR + notInValue); - if (dictId >= 0) { - notInDictIds.add(dictId); - } - } - } - for (int dictId = dictIds[0]; dictId < dictIds[1]; dictId++) { - if (notInDictIds.contains(dictId)) { - continue; + IntOpenHashSet notInDictIds = new IntOpenHashSet(); + for (String value : values) { + buffer.setLength(pos); + buffer.append(value); + int dictId = _dictionary.indexOf(buffer.toString()); + if (dictId >= 0) { + notInDictIds.add(dictId); } - - if (result == null) { - result = _invertedIndex.getDocIds(dictId); - } else { + } + ImmutableRoaringBitmap result = EMPTY_BITMAP; + for (int dictId = dictIdRange[0]; dictId < dictIdRange[1]; dictId++) { + if (!notInDictIds.contains(dictId)) { result = or(result, _invertedIndex.getDocIds(dictId)); } } + return result; } - - return filter(result, matchingDocIds); } case IS_NOT_NULL: case IS_NULL: { - ImmutableRoaringBitmap result = null; int dictId = _dictionary.indexOf(key); if (dictId >= 0) { - result = _invertedIndex.getDocIds(dictId); + return _invertedIndex.getDocIds(dictId); + } else { + return EMPTY_BITMAP; } - - return filter(result, matchingDocIds); } case REGEXP_LIKE: { + int[] dictIds = getDictIdRangeForKey(key); + int minDictId = dictIds[0]; + if (minDictId < 0) { + return EMPTY_BITMAP; + } Pattern pattern = ((RegexpLikePredicate) predicate).getPattern(); Matcher matcher = pattern.matcher(""); - int[] dictIds = getDictIdRangeForKey(key); - - ImmutableRoaringBitmap result = null; - byte[] dictBuffer = dictIds[0] < dictIds[1] ? _dictionary.getBuffer() : null; + ImmutableRoaringBitmap result = EMPTY_BITMAP; + byte[] dictBuffer = _dictionary.getBuffer(); StringBuilder value = new StringBuilder(); - + int valueStart = key.length() + 1; for (int dictId = dictIds[0]; dictId < dictIds[1]; dictId++) { - String stringValue = _dictionary.getStringValue(dictId, dictBuffer); + String keyValue = _dictionary.getStringValue(dictId, dictBuffer); value.setLength(0); - value.append(stringValue, key.length() + 1, stringValue.length()); - + value.append(keyValue, valueStart, keyValue.length()); if (matcher.reset(value).matches()) { - if (result == null) { - result = _invertedIndex.getDocIds(dictId); - } else { - result = or(result, _invertedIndex.getDocIds(dictId)); - } + result = or(result, _invertedIndex.getDocIds(dictId)); } } - - return filter(result, matchingDocIds); + return result; } case RANGE: { + int[] dictIds = getDictIdRangeForKey(key); + int minDictId = dictIds[0]; + if (minDictId < 0) { + return EMPTY_BITMAP; + } RangePredicate rangePredicate = (RangePredicate) predicate; FieldSpec.DataType rangeDataType = rangePredicate.getRangeDataType(); // Simplify to only support numeric and string types @@ -470,20 +438,17 @@ private ImmutableRoaringBitmap getMatchingFlattenedDocIds(Predicate predicate) { } else { rangeDataType = FieldSpec.DataType.STRING; } - boolean lowerUnbounded = rangePredicate.getLowerBound().equals(RangePredicate.UNBOUNDED); boolean upperUnbounded = rangePredicate.getUpperBound().equals(RangePredicate.UNBOUNDED); boolean lowerInclusive = lowerUnbounded || rangePredicate.isLowerInclusive(); boolean upperInclusive = upperUnbounded || rangePredicate.isUpperInclusive(); Object lowerBound = lowerUnbounded ? null : rangeDataType.convert(rangePredicate.getLowerBound()); Object upperBound = upperUnbounded ? null : rangeDataType.convert(rangePredicate.getUpperBound()); - - int[] dictIds = getDictIdRangeForKey(key); - ImmutableRoaringBitmap result = null; - byte[] dictBuffer = dictIds[0] < dictIds[1] ? _dictionary.getBuffer() : null; - + ImmutableRoaringBitmap result = EMPTY_BITMAP; + byte[] dictBuffer = _dictionary.getBuffer(); + int valueStart = key.length() + 1; for (int dictId = dictIds[0]; dictId < dictIds[1]; dictId++) { - String value = _dictionary.getStringValue(dictId, dictBuffer).substring(key.length() + 1); + String value = _dictionary.getStringValue(dictId, dictBuffer).substring(valueStart); Object valueObj = rangeDataType.convert(value); boolean lowerCompareResult = lowerUnbounded || (lowerInclusive ? rangeDataType.compare(valueObj, lowerBound) >= 0 @@ -491,17 +456,11 @@ private ImmutableRoaringBitmap getMatchingFlattenedDocIds(Predicate predicate) { boolean upperCompareResult = upperUnbounded || (upperInclusive ? rangeDataType.compare(valueObj, upperBound) <= 0 : rangeDataType.compare(valueObj, upperBound) < 0); - if (lowerCompareResult && upperCompareResult) { - if (result == null) { - result = _invertedIndex.getDocIds(dictId); - } else { - result = or(result, _invertedIndex.getDocIds(dictId)); - } + result = or(result, _invertedIndex.getDocIds(dictId)); } } - - return filter(result, matchingDocIds); + return result; } default: @@ -591,8 +550,7 @@ public Map getMatchingFlattenedDocsMap(String jsonPathKey } @Override - public String[][] getValuesMV(int[] docIds, int length, - Map valueToMatchingFlattenedDocs) { + public String[][] getValuesMV(int[] docIds, int length, Map valueToMatchingFlattenedDocs) { String[][] result = new String[length][]; List>> docIdToFlattenedDocIdsAndValues = new ArrayList<>(); for (int i = 0; i < length; i++) { @@ -714,7 +672,7 @@ private Pair getKeyAndFlattenedDocIds(String key // "[0]"=1 -> ".$index"='0' && "."='1' // ".foo[1].bar"='abc' -> ".foo.$index"=1 && ".foo..bar"='abc' String searchKey = - leftPart + JsonUtils.ARRAY_INDEX_KEY + BaseJsonIndexCreator.KEY_VALUE_SEPARATOR + arrayIndex; + leftPart + JsonUtils.ARRAY_INDEX_KEY + BaseJsonIndexCreator.KEY_VALUE_SEPARATOR + arrayIndex; int dictId = _dictionary.indexOf(searchKey); if (dictId >= 0) { ImmutableRoaringBitmap docIds = _invertedIndex.getDocIds(dictId); @@ -746,7 +704,7 @@ private Pair getKeyAndFlattenedDocIds(String key if (!arrayIndex.equals(JsonUtils.WILDCARD)) { // "foo[1].bar"='abc' -> "foo.$index"=1 && "foo.bar"='abc' String searchKey = - leftPart + JsonUtils.ARRAY_INDEX_KEY + BaseJsonIndexCreator.KEY_VALUE_SEPARATOR + arrayIndex; + leftPart + JsonUtils.ARRAY_INDEX_KEY + BaseJsonIndexCreator.KEY_VALUE_SEPARATOR + arrayIndex; int dictId = _dictionary.indexOf(searchKey); if (dictId >= 0) { ImmutableRoaringBitmap docIds = _invertedIndex.getDocIds(dictId); @@ -766,11 +724,6 @@ private Pair getKeyAndFlattenedDocIds(String key return Pair.of(key, matchingDocIds); } - private PeekableIntIterator intersect(MutableRoaringBitmap a, ImmutableRoaringBitmap b) { - a.and(b); - return a.getIntIterator(); - } - @Override public void close() { // NOTE: DO NOT close the PinotDataBuffer here because it is tracked by the caller and might be reused later. The From b1a98eeac8361214f50f8cb3cf87b29435572883 Mon Sep 17 00:00:00 2001 From: Hongkun Xu Date: Tue, 29 Jul 2025 06:47:02 +0800 Subject: [PATCH 031/167] Add auth for revert replace segment (#16434) --- .../utils/FileUploadDownloadClient.java | 24 ++++++++++++++++++- .../local/utils/ConsistentDataPushUtils.java | 3 ++- 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/FileUploadDownloadClient.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/FileUploadDownloadClient.java index fab48220fd25..543f6c80f74d 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/FileUploadDownloadClient.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/FileUploadDownloadClient.java @@ -517,6 +517,13 @@ private static ClassicHttpRequest getRevertReplaceSegmentRequest(URI uri) { return requestBuilder.build(); } + private static ClassicHttpRequest getRevertReplaceSegmentRequest(URI uri, @Nullable AuthProvider authProvider) { + ClassicRequestBuilder requestBuilder = ClassicRequestBuilder.post(uri).setVersion(HttpVersion.HTTP_1_1) + .setHeader(HttpHeaders.CONTENT_TYPE, HttpClient.JSON_CONTENT_TYPE); + AuthProviderUtils.toRequestHeaders(authProvider).forEach(requestBuilder::addHeader); + return requestBuilder.build(); + } + private static ClassicHttpRequest getSegmentCompletionProtocolRequest(URI uri, @Nullable List
headers, @Nullable List parameters) { ClassicRequestBuilder requestBuilder = ClassicRequestBuilder.get(uri).setVersion(HttpVersion.HTTP_1_1); @@ -1174,7 +1181,22 @@ public SimpleHttpResponse endReplaceSegments(URI uri, int socketTimeoutMs, */ public SimpleHttpResponse revertReplaceSegments(URI uri) throws IOException, HttpErrorStatusException { - return HttpClient.wrapAndThrowHttpException(_httpClient.sendRequest(getRevertReplaceSegmentRequest(uri))); + return revertReplaceSegments(uri, null); + } + + /** + * Revert replace segments with default settings. + * + * @param uri URI + * @param authProvider auth provider + * @return Response + * @throws IOException + * @throws HttpErrorStatusException + */ + public SimpleHttpResponse revertReplaceSegments(URI uri, @Nullable AuthProvider authProvider) + throws IOException, HttpErrorStatusException { + return HttpClient.wrapAndThrowHttpException(_httpClient.sendRequest( + getRevertReplaceSegmentRequest(uri, authProvider))); } /** diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/ConsistentDataPushUtils.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/ConsistentDataPushUtils.java index 798d839c9eb1..f7a07083fc87 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/ConsistentDataPushUtils.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/ConsistentDataPushUtils.java @@ -197,12 +197,13 @@ public static void handleUploadException(SegmentGenerationJobSpec spec, Map entry : uriToLineageEntryIdMap.entrySet()) { String segmentLineageEntryId = entry.getValue(); try { URI uri = FileUploadDownloadClient.getRevertReplaceSegmentsURI(entry.getKey(), rawTableName, TableType.OFFLINE.name(), segmentLineageEntryId, true); - SimpleHttpResponse response = FILE_UPLOAD_DOWNLOAD_CLIENT.revertReplaceSegments(uri); + SimpleHttpResponse response = FILE_UPLOAD_DOWNLOAD_CLIENT.revertReplaceSegments(uri, authProvider); LOGGER.info("Got response {}: {} while sending revert replace segment request for table: {}, uploadURI: {}", response.getStatusCode(), response.getResponse(), rawTableName, entry.getKey()); } catch (URISyntaxException | HttpErrorStatusException | IOException e) { From 3121f8fc56f28f3e76c3c14d7bc86b29da91b96d Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Mon, 28 Jul 2025 19:03:45 -0700 Subject: [PATCH 032/167] Add state and table filtering support to task counts API (#16433) * Add state and table filtering support to task counts API - Enhanced /tasks/{taskType}/taskcounts endpoint with optional query parameters: * state: Filter by single or multiple comma-separated task states (waiting, running, error, completed, dropped, timedOut, aborted, unknown, total) * table: Filter by table name to show only tasks with subtasks for specific tables - Features: * Multiple state filtering: ?state=running,error,waiting * Table-specific filtering: ?table=myTable_OFFLINE * Combined filtering: ?state=running&table=myTable_OFFLINE * Database-aware table name translation * Comprehensive input validation with clear error messages * Backward compatible - existing API usage unchanged - Implementation: * New overloaded getTaskCounts methods in PinotHelixTaskResourceManager * Efficient filtering logic with early validation * Robust error handling for edge cases * Helper methods for state and table filtering logic - Tests: * Added comprehensive test suite covering all scenarios * State filtering tests (single, multiple, invalid states) * Table filtering tests (table-only, combined with states) * Edge case handling (null values, exceptions, no matches) * 6 new test methods with proper mocking This enhancement enables precise monitoring of Pinot tasks by allowing operators to filter task counts by both execution state and target table, significantly improving operational visibility. * Update task counts API to use Helix TaskState enum for state filtering - Changed state parameter to use official Helix TaskState enum values: NOT_STARTED, IN_PROGRESS, STOPPED, STOPPING, FAILED, COMPLETED, ABORTED, TIMED_OUT, TIMING_OUT, FAILING - Updated API documentation with correct TaskState values and examples - Enhanced state validation using TaskState.valueOf() for robust enum validation - Updated filtering logic to work with actual TaskState enum instead of aggregated counts - Added validateAndParseTaskState() method for proper state validation with clear error messages - Updated all test methods to use correct TaskState enum values - Fixed all checkstyle violations and line length issues This ensures the API uses the standard Helix TaskState semantics for consistent task monitoring across the Pinot ecosystem. * Fix PinotHelixTaskResourceManagerTest workflow context mocking - Fixed failing tests for TaskState enum filtering functionality - Updated workflow context mocking to use getJobStates() method instead of individual getJobState() calls - Added proper WorkflowConfig and JobDag mocking for getTaskStates() method dependencies - Fixed import organization and removed duplicate imports - All TaskState filtering tests now pass: testGetTaskCountsWithSingleStateFilter, testGetTaskCountsWithMultipleStatesFilter, testGetTaskCountsWithStateAndTableFilter The tests now properly mock the Helix workflow dependencies required by the updated TaskState enum implementation, ensuring correct state filtering behavior. * Optimize getTaskCounts performance by applying filters before collecting TaskCount - Move expensive getTaskCount() operation to only happen after state and table filtering - Apply state filtering first, then table filtering, then collect TaskCount - Significantly improves performance when many tasks are filtered out - Addresses review feedback from krishan1390 on PR #16433 * Address PR review comments: use StringUtils.isNotEmpty() and add @Nullable annotations - Replace null checks with StringUtils.isNotEmpty() for better string validation in PinotTaskRestletResource - Add @Nullable annotations to state and tableNameWithType parameters in PinotHelixTaskResourceManager.getTaskCounts() - Update state and table filtering logic to use StringUtils.isNotEmpty() for consistency - Improve method signature formatting for better readability Addresses feedback from Jackie-Jiang on PR #16433 --- .../resources/PinotTaskRestletResource.java | 38 +- .../minion/PinotHelixTaskResourceManager.java | 115 +++++ .../PinotTaskRestletResourceTest.java | 1 - .../PinotHelixTaskResourceManagerTest.java | 405 +++++++++++++++++- 4 files changed, 543 insertions(+), 16 deletions(-) diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTaskRestletResource.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTaskRestletResource.java index d76ea4f3a254..90e6732dc968 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTaskRestletResource.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTaskRestletResource.java @@ -130,15 +130,18 @@ *
  • DELETE '/tasks/{taskType}': Delete all tasks (as well as the task queue) for the given task type
    • * */ -@Api(tags = Constants.TASK_TAG, authorizations = {@Authorization(value = SWAGGER_AUTHORIZATION_KEY), - @Authorization(value = DATABASE)}) +@Api(tags = Constants.TASK_TAG, authorizations = { + @Authorization(value = SWAGGER_AUTHORIZATION_KEY), + @Authorization(value = DATABASE) +}) @SwaggerDefinition(securityDefinition = @SecurityDefinition(apiKeyAuthDefinitions = { @ApiKeyAuthDefinition(name = HttpHeaders.AUTHORIZATION, in = ApiKeyAuthDefinition.ApiKeyLocation.HEADER, key = SWAGGER_AUTHORIZATION_KEY, description = "The format of the key is ```\"Basic \" or \"Bearer \"```"), @ApiKeyAuthDefinition(name = DATABASE, in = ApiKeyAuthDefinition.ApiKeyLocation.HEADER, key = DATABASE, description = "Database context passed through http header. If no context is provided 'default' database " - + "context will be considered.")})) + + "context will be considered.") +})) @Path("/") public class PinotTaskRestletResource { public static final Logger LOGGER = LoggerFactory.getLogger(PinotTaskRestletResource.class); @@ -267,8 +270,21 @@ public SuccessResponse deleteTaskMetadataByTable( @Produces(MediaType.APPLICATION_JSON) @ApiOperation("Fetch count of sub-tasks for each of the tasks for the given task type") public Map getTaskCounts( - @ApiParam(value = "Task type", required = true) @PathParam("taskType") String taskType) { - return _pinotHelixTaskResourceManager.getTaskCounts(taskType); + @ApiParam(value = "Task type", required = true) @PathParam("taskType") String taskType, + @ApiParam(value = "Task state(s) to filter by. Can be single state or comma-separated multiple states " + + "(NOT_STARTED, IN_PROGRESS, STOPPED, STOPPING, FAILED, COMPLETED, ABORTED, TIMED_OUT, TIMING_OUT, " + + "FAILING). Example: 'IN_PROGRESS' or 'IN_PROGRESS,FAILED,STOPPING'") + @QueryParam("state") @Nullable String state, + @ApiParam(value = "Table name with type (e.g., 'myTable_OFFLINE') to filter tasks by table. " + + "Only tasks that have subtasks for this table will be returned.") + @QueryParam("table") @Nullable String table, @Context HttpHeaders headers) { + String tableNameWithType = table != null ? DatabaseUtils.translateTableName(table, headers) : null; + + if (StringUtils.isNotEmpty(state) || StringUtils.isNotEmpty(tableNameWithType)) { + return _pinotHelixTaskResourceManager.getTaskCounts(taskType, state, tableNameWithType); + } else { + return _pinotHelixTaskResourceManager.getTaskCounts(taskType); + } } @GET @@ -293,7 +309,7 @@ public Map getTasksDebugInf public Map getTasksDebugInfo( @ApiParam(value = "Task type", required = true) @PathParam("taskType") String taskType, @ApiParam(value = "Table name with type", required = true) @PathParam("tableNameWithType") - String tableNameWithType, + String tableNameWithType, @ApiParam(value = "verbosity (Prints information for all the tasks for the given task type and table." + "By default, only prints subtask details for running and error tasks. " + "Value of > 0 prints subtask details for all tasks)") @@ -310,9 +326,9 @@ public Map getTasksDebugInf public String getTaskGenerationDebugInto( @ApiParam(value = "Task type", required = true) @PathParam("taskType") String taskType, @ApiParam(value = "Table name with type", required = true) @PathParam("tableNameWithType") - String tableNameWithType, + String tableNameWithType, @ApiParam(value = "Whether to only lookup local cache for logs", defaultValue = "false") @QueryParam("localOnly") - boolean localOnly, @Context HttpHeaders httpHeaders) + boolean localOnly, @Context HttpHeaders httpHeaders) throws JsonProcessingException { tableNameWithType = DatabaseUtils.translateTableName(tableNameWithType, httpHeaders); if (localOnly) { @@ -430,7 +446,7 @@ public Map getTaskConfig( public Map getSubtaskConfigs( @ApiParam(value = "Task name", required = true) @PathParam("taskName") String taskName, @ApiParam(value = "Sub task names separated by comma") @QueryParam("subtaskNames") @Nullable - String subtaskNames) { + String subtaskNames) { return _pinotHelixTaskResourceManager.getSubtaskConfigs(taskName, subtaskNames); } @@ -442,7 +458,7 @@ public Map getSubtaskConfigs( public String getSubtaskProgress(@Context HttpHeaders httpHeaders, @ApiParam(value = "Task name", required = true) @PathParam("taskName") String taskName, @ApiParam(value = "Sub task names separated by comma") @QueryParam("subtaskNames") @Nullable - String subtaskNames) { + String subtaskNames) { // Relying on original schema that was used to query the controller String scheme = _uriInfo.getRequestUri().getScheme(); List workers = _pinotHelixResourceManager.getAllMinionInstanceConfigs(); @@ -482,7 +498,7 @@ public String getSubtaskOnWorkerProgress(@Context HttpHeaders httpHeaders, @ApiParam(value = "Subtask state (UNKNOWN,IN_PROGRESS,SUCCEEDED,CANCELLED,ERROR)", required = true) @QueryParam("subTaskState") String subTaskState, @ApiParam(value = "Minion worker IDs separated by comma") @QueryParam("minionWorkerIds") @Nullable - String minionWorkerIds) { + String minionWorkerIds) { Set selectedMinionWorkers = new HashSet<>(); if (StringUtils.isNotEmpty(minionWorkerIds)) { selectedMinionWorkers.addAll( diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/minion/PinotHelixTaskResourceManager.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/minion/PinotHelixTaskResourceManager.java index 6f1fb4c0a22a..5f86db52b0a4 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/minion/PinotHelixTaskResourceManager.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/minion/PinotHelixTaskResourceManager.java @@ -776,6 +776,121 @@ public synchronized Map getTaskCounts(String taskType) { return taskCounts; } + /** + * Fetch count of sub-tasks for each of the tasks for the given taskType, filtered by state. + * + * @param taskType Pinot taskType / Helix JobQueue + * @param state State(s) to filter by. Can be single state or comma-separated multiple states + * (waiting, running, error, completed, dropped, timedOut, aborted, unknown, total) + * @return Map of Pinot Task Name to TaskCount containing only tasks that have > 0 count for any of the + * specified states + */ + public synchronized Map getTaskCounts(String taskType, String state) { + return getTaskCounts(taskType, state, null); + } + + /** + * Fetch count of sub-tasks for each of the tasks for the given taskType, filtered by state and/or table. + * + * @param taskType Pinot taskType / Helix JobQueue + * @param state State(s) to filter by. Can be single state or comma-separated multiple states + * (NOT_STARTED, IN_PROGRESS, STOPPED, STOPPING, FAILED, COMPLETED, ABORTED, TIMED_OUT, + * TIMING_OUT, FAILING). Can be null to skip state filtering. + * @param tableNameWithType Table name with type to filter by. Only tasks that have subtasks for this table + * will be returned. Can be null to skip table filtering. + * @return Map of Pinot Task Name to TaskCount containing only tasks that match the specified filters + */ + public synchronized Map getTaskCounts(String taskType, @Nullable String state, + @Nullable String tableNameWithType) { + Set tasks = getTasks(taskType); + if (tasks == null) { + return Collections.emptyMap(); + } + + // Parse and validate comma-separated states if provided + Set requestedStates = null; + if (StringUtils.isNotEmpty(state)) { + String[] stateArray = state.trim().split(","); + requestedStates = new HashSet<>(); + for (String s : stateArray) { + String normalizedState = s.trim().toUpperCase(); + // Validate each state upfront + TaskState taskState = validateAndParseTaskState(normalizedState); + requestedStates.add(taskState); + } + } + + // Get all task states if we need to filter by state + Map taskStates = null; + if (requestedStates != null) { + taskStates = getTaskStates(taskType); + } + + Map taskCounts = new TreeMap<>(); + for (String taskName : tasks) { + // Apply state filtering first (less expensive) + if (requestedStates != null) { + TaskState currentTaskState = taskStates.get(taskName); + if (currentTaskState == null || !requestedStates.contains(currentTaskState)) { + continue; + } + } + + // Apply table filtering next (also less expensive than getting task count) + if (StringUtils.isNotEmpty(tableNameWithType) && !hasTasksForTable(taskName, tableNameWithType)) { + continue; + } + + // Only collect TaskCount after passing all filters (expensive operation) + TaskCount taskCount = getTaskCount(taskName); + taskCounts.put(taskName, taskCount); + } + return taskCounts; + } + + /** + * Validates and parses a task state string into TaskState enum. + * + * @param state State string to validate (should be uppercase) + * @throws IllegalArgumentException if the state is invalid + * @return TaskState enum value + */ + private TaskState validateAndParseTaskState(String state) { + try { + return TaskState.valueOf(state); + } catch (IllegalArgumentException e) { + throw new IllegalArgumentException("Invalid state: " + state + ". Valid states are: " + + Arrays.toString(TaskState.values())); + } + } + + /** + * Helper method to check if a task has any subtasks for the specified table. + * + * @param taskName Task name to check + * @param tableNameWithType Table name with type to check for + * @return true if the task has subtasks for the specified table + */ + private boolean hasTasksForTable(String taskName, String tableNameWithType) { + try { + // Get all subtask configs for this task + List subtaskConfigs = getSubtaskConfigs(taskName); + + // Check if any subtask is for the specified table + for (PinotTaskConfig taskConfig : subtaskConfigs) { + String taskTableName = taskConfig.getTableName(); + if (taskTableName != null && taskTableName.equals(tableNameWithType)) { + return true; + } + } + return false; + } catch (Exception e) { + // If we can't get the subtask configs, assume no match + LOGGER.warn("Failed to get subtask configs for task: {}", taskName, e); + return false; + } + } + /** * Given a taskType, helper method to debug all the HelixJobs for the taskType. * For each of the HelixJobs, collects status of the (sub)tasks in the taskbatch. diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/PinotTaskRestletResourceTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/PinotTaskRestletResourceTest.java index 764b1d53d16b..3e1fac572768 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/PinotTaskRestletResourceTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/PinotTaskRestletResourceTest.java @@ -123,7 +123,6 @@ private InstanceConfig createInstanceConfig(String instanceId, String hostName, return instanceConfig; } - @Test public void testGetSubtaskWithGivenStateProgressWithException() throws JsonProcessingException { diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/minion/PinotHelixTaskResourceManagerTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/minion/PinotHelixTaskResourceManagerTest.java index ef11b4f30202..cee3c246f460 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/minion/PinotHelixTaskResourceManagerTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/minion/PinotHelixTaskResourceManagerTest.java @@ -21,6 +21,7 @@ import com.fasterxml.jackson.core.JsonProcessingException; import com.google.common.collect.ImmutableMap; import java.io.IOException; +import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; @@ -31,10 +32,12 @@ import org.apache.commons.lang3.StringUtils; import org.apache.helix.task.JobConfig; import org.apache.helix.task.JobContext; +import org.apache.helix.task.JobDag; import org.apache.helix.task.TaskConfig; import org.apache.helix.task.TaskDriver; import org.apache.helix.task.TaskPartitionState; import org.apache.helix.task.TaskState; +import org.apache.helix.task.WorkflowConfig; import org.apache.helix.task.WorkflowContext; import org.apache.pinot.controller.helix.core.PinotHelixResourceManager; import org.apache.pinot.controller.util.CompletionServiceHelper; @@ -52,10 +55,7 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; -import static org.testng.Assert.assertEquals; -import static org.testng.Assert.assertNotNull; -import static org.testng.Assert.assertNull; -import static org.testng.Assert.assertTrue; +import static org.testng.Assert.*; public class PinotHelixTaskResourceManagerTest { @@ -664,4 +664,401 @@ public void testGetTasksDebugInfoByTableWithVerbosity() { assertEquals(subtaskInfos.size(), 1); // Completed tasks should be included assertEquals(subtaskInfos.get(0).getState(), TaskPartitionState.COMPLETED); } + + @Test + public void testGetTaskCountsWithSingleStateFilter() { + String taskType = "TestTask"; + String taskName1 = "Task_TestTask_12345"; + String taskName2 = "Task_TestTask_67890"; + + TaskDriver taskDriver = mock(TaskDriver.class); + PinotHelixResourceManager pinotHelixResourceManager = mock(PinotHelixResourceManager.class); + PinotHelixTaskResourceManager mgr = new PinotHelixTaskResourceManager(pinotHelixResourceManager, taskDriver); + + // Mock getTasks to return our test task names + Set tasks = new HashSet<>(); + tasks.add(taskName1); + tasks.add(taskName2); + PinotHelixTaskResourceManager spyMgr = Mockito.spy(mgr); + when(spyMgr.getTasks(taskType)).thenReturn(tasks); + + // Mock workflow-level components for getTaskStates + String helixJobQueueName = "TaskQueue_" + taskType; + WorkflowConfig workflowConfig = mock(WorkflowConfig.class); + when(taskDriver.getWorkflowConfig(helixJobQueueName)).thenReturn(workflowConfig); + + JobDag jobDag = mock(JobDag.class); + Set helixJobs = new HashSet<>(); + helixJobs.add(PinotHelixTaskResourceManager.getHelixJobName(taskName1)); + helixJobs.add(PinotHelixTaskResourceManager.getHelixJobName(taskName2)); + when(jobDag.getAllNodes()).thenReturn(helixJobs); + when(workflowConfig.getJobDag()).thenReturn(jobDag); + + WorkflowContext workflowContext = mock(WorkflowContext.class); + when(taskDriver.getWorkflowContext(helixJobQueueName)).thenReturn(workflowContext); + Map jobStatesMap = new HashMap<>(); + jobStatesMap.put(PinotHelixTaskResourceManager.getHelixJobName(taskName1), TaskState.IN_PROGRESS); + jobStatesMap.put(PinotHelixTaskResourceManager.getHelixJobName(taskName2), TaskState.COMPLETED); + when(workflowContext.getJobStates()).thenReturn(jobStatesMap); + + // Mock JobConfig and JobContext for both tasks + mockTaskJobConfigAndContext(taskDriver, taskName1, TaskPartitionState.RUNNING); + mockTaskJobConfigAndContext(taskDriver, taskName2, TaskPartitionState.COMPLETED); + + // Test filter by "IN_PROGRESS" - should only return taskName1 + Map inProgressTasks = + spyMgr.getTaskCounts(taskType, "IN_PROGRESS", null); + assertEquals(inProgressTasks.size(), 1); + assertTrue(inProgressTasks.containsKey(taskName1)); + assertFalse(inProgressTasks.containsKey(taskName2)); + + // Test filter by "COMPLETED" - should only return taskName2 + Map completedTasks = + spyMgr.getTaskCounts(taskType, "COMPLETED", null); + assertEquals(completedTasks.size(), 1); + assertFalse(completedTasks.containsKey(taskName1)); + assertTrue(completedTasks.containsKey(taskName2)); + + // Test filter by "FAILED" - should return no tasks + Map failedTasks = + spyMgr.getTaskCounts(taskType, "FAILED", null); + assertEquals(failedTasks.size(), 0); + } + + @Test + public void testGetTaskCountsWithMultipleStatesFilter() { + String taskType = "TestTask"; + String taskName1 = "Task_TestTask_12345"; + String taskName2 = "Task_TestTask_67890"; + String taskName3 = "Task_TestTask_11111"; + + TaskDriver taskDriver = mock(TaskDriver.class); + PinotHelixResourceManager pinotHelixResourceManager = mock(PinotHelixResourceManager.class); + PinotHelixTaskResourceManager mgr = new PinotHelixTaskResourceManager(pinotHelixResourceManager, taskDriver); + + // Mock getTasks to return our test task names + Set tasks = new HashSet<>(); + tasks.add(taskName1); + tasks.add(taskName2); + tasks.add(taskName3); + PinotHelixTaskResourceManager spyMgr = Mockito.spy(mgr); + when(spyMgr.getTasks(taskType)).thenReturn(tasks); + + // Mock workflow-level components for getTaskStates + String helixJobQueueName = "TaskQueue_" + taskType; + WorkflowConfig workflowConfig = mock(WorkflowConfig.class); + when(taskDriver.getWorkflowConfig(helixJobQueueName)).thenReturn(workflowConfig); + + JobDag jobDag = mock(JobDag.class); + Set helixJobs = new HashSet<>(); + helixJobs.add(PinotHelixTaskResourceManager.getHelixJobName(taskName1)); + helixJobs.add(PinotHelixTaskResourceManager.getHelixJobName(taskName2)); + helixJobs.add(PinotHelixTaskResourceManager.getHelixJobName(taskName3)); + when(jobDag.getAllNodes()).thenReturn(helixJobs); + when(workflowConfig.getJobDag()).thenReturn(jobDag); + + WorkflowContext workflowContext = mock(WorkflowContext.class); + when(taskDriver.getWorkflowContext(helixJobQueueName)).thenReturn(workflowContext); + Map jobStatesMap = new HashMap<>(); + jobStatesMap.put(PinotHelixTaskResourceManager.getHelixJobName(taskName1), TaskState.IN_PROGRESS); + jobStatesMap.put(PinotHelixTaskResourceManager.getHelixJobName(taskName2), TaskState.FAILED); + jobStatesMap.put(PinotHelixTaskResourceManager.getHelixJobName(taskName3), TaskState.COMPLETED); + when(workflowContext.getJobStates()).thenReturn(jobStatesMap); + + // Mock JobConfig and JobContext for all tasks + mockTaskJobConfigAndContext(taskDriver, taskName1, TaskPartitionState.RUNNING); + mockTaskJobConfigAndContext(taskDriver, taskName2, TaskPartitionState.TASK_ERROR); + mockTaskJobConfigAndContext(taskDriver, taskName3, TaskPartitionState.COMPLETED); + + // Test filter by "IN_PROGRESS,FAILED" - should return taskName1 and taskName2 + Map inProgressOrFailedTasks = + spyMgr.getTaskCounts(taskType, "IN_PROGRESS,FAILED", null); + assertEquals(inProgressOrFailedTasks.size(), 2); + assertTrue(inProgressOrFailedTasks.containsKey(taskName1)); + assertTrue(inProgressOrFailedTasks.containsKey(taskName2)); + assertFalse(inProgressOrFailedTasks.containsKey(taskName3)); + + // Test filter by "COMPLETED,IN_PROGRESS" - should return taskName1 and taskName3 + Map completedOrInProgressTasks = + spyMgr.getTaskCounts(taskType, "COMPLETED,IN_PROGRESS", null); + assertEquals(completedOrInProgressTasks.size(), 2); + assertTrue(completedOrInProgressTasks.containsKey(taskName1)); + assertFalse(completedOrInProgressTasks.containsKey(taskName2)); + assertTrue(completedOrInProgressTasks.containsKey(taskName3)); + + // Test filter by "NOT_STARTED,STOPPED" - should return no tasks + Map notStartedOrStoppedTasks = + spyMgr.getTaskCounts(taskType, "NOT_STARTED,STOPPED", null); + assertEquals(notStartedOrStoppedTasks.size(), 0); + + // Test filter with spaces "IN_PROGRESS, FAILED, COMPLETED" - should return all three + Map allTasks = + spyMgr.getTaskCounts(taskType, "IN_PROGRESS, FAILED, COMPLETED", null); + assertEquals(allTasks.size(), 3); + assertTrue(allTasks.containsKey(taskName1)); + assertTrue(allTasks.containsKey(taskName2)); + assertTrue(allTasks.containsKey(taskName3)); + } + + @Test + public void testGetTaskCountsWithInvalidState() { + String taskType = "TestTask"; + String taskName1 = "Task_TestTask_12345"; + + TaskDriver taskDriver = mock(TaskDriver.class); + PinotHelixResourceManager pinotHelixResourceManager = mock(PinotHelixResourceManager.class); + PinotHelixTaskResourceManager mgr = new PinotHelixTaskResourceManager(pinotHelixResourceManager, taskDriver); + + // Mock getTasks to return a task + Set tasks = new HashSet<>(); + tasks.add(taskName1); + PinotHelixTaskResourceManager spyMgr = Mockito.spy(mgr); + when(spyMgr.getTasks(taskType)).thenReturn(tasks); + + // Mock JobConfig and JobContext for the task + mockTaskJobConfigAndContext(taskDriver, taskName1, TaskPartitionState.RUNNING); + + // Test with invalid single state + try { + spyMgr.getTaskCounts(taskType, "INVALID_STATE", null); + fail("Expected IllegalArgumentException for invalid state"); + } catch (IllegalArgumentException e) { + assertTrue(e.getMessage().contains("Invalid state: INVALID_STATE")); + } + + // Test with mixed valid and invalid states + try { + spyMgr.getTaskCounts(taskType, "IN_PROGRESS,INVALID_STATE,COMPLETED", null); + fail("Expected IllegalArgumentException for invalid state in multiple states"); + } catch (IllegalArgumentException e) { + assertTrue(e.getMessage().contains("Invalid state: INVALID_STATE")); + } + } + + @Test + public void testGetTaskCountsWithTableFilter() { + String taskType = "TestTask"; + String taskName1 = "Task_TestTask_12345"; + String taskName2 = "Task_TestTask_67890"; + String table1 = "table1_OFFLINE"; + String table2 = "table2_OFFLINE"; + + TaskDriver taskDriver = mock(TaskDriver.class); + PinotHelixResourceManager pinotHelixResourceManager = mock(PinotHelixResourceManager.class); + PinotHelixTaskResourceManager mgr = new PinotHelixTaskResourceManager(pinotHelixResourceManager, taskDriver); + + // Mock getTasks to return our test task names + Set tasks = new HashSet<>(); + tasks.add(taskName1); + tasks.add(taskName2); + PinotHelixTaskResourceManager spyMgr = Mockito.spy(mgr); + when(spyMgr.getTasks(taskType)).thenReturn(tasks); + + // Mock JobConfig and JobContext for both tasks + mockTaskJobConfigAndContext(taskDriver, taskName1, TaskPartitionState.RUNNING); + mockTaskJobConfigAndContext(taskDriver, taskName2, TaskPartitionState.COMPLETED); + + // Mock subtask configs - taskName1 has subtasks for table1, taskName2 has subtasks for table2 + List subtaskConfigs1 = new ArrayList<>(); + PinotTaskConfig taskConfig1 = mock(PinotTaskConfig.class); + when(taskConfig1.getTableName()).thenReturn(table1); + subtaskConfigs1.add(taskConfig1); + when(spyMgr.getSubtaskConfigs(taskName1)).thenReturn(subtaskConfigs1); + + List subtaskConfigs2 = new ArrayList<>(); + PinotTaskConfig taskConfig2 = mock(PinotTaskConfig.class); + when(taskConfig2.getTableName()).thenReturn(table2); + subtaskConfigs2.add(taskConfig2); + when(spyMgr.getSubtaskConfigs(taskName2)).thenReturn(subtaskConfigs2); + + // Test filter by table1 - should only return taskName1 + Map table1Tasks = + spyMgr.getTaskCounts(taskType, null, table1); + assertEquals(table1Tasks.size(), 1); + assertTrue(table1Tasks.containsKey(taskName1)); + assertFalse(table1Tasks.containsKey(taskName2)); + + // Test filter by table2 - should only return taskName2 + Map table2Tasks = + spyMgr.getTaskCounts(taskType, null, table2); + assertEquals(table2Tasks.size(), 1); + assertFalse(table2Tasks.containsKey(taskName1)); + assertTrue(table2Tasks.containsKey(taskName2)); + + // Test filter by non-existent table - should return no tasks + Map noTableTasks = + spyMgr.getTaskCounts(taskType, null, "nonexistent_OFFLINE"); + assertEquals(noTableTasks.size(), 0); + } + + @Test + public void testGetTaskCountsWithStateAndTableFilter() { + String taskType = "TestTask"; + String taskName1 = "Task_TestTask_12345"; + String taskName2 = "Task_TestTask_67890"; + String taskName3 = "Task_TestTask_11111"; + String table1 = "table1_OFFLINE"; + String table2 = "table2_OFFLINE"; + + TaskDriver taskDriver = mock(TaskDriver.class); + PinotHelixResourceManager pinotHelixResourceManager = mock(PinotHelixResourceManager.class); + PinotHelixTaskResourceManager mgr = new PinotHelixTaskResourceManager(pinotHelixResourceManager, taskDriver); + + // Mock getTasks to return our test task names + Set tasks = new HashSet<>(); + tasks.add(taskName1); + tasks.add(taskName2); + tasks.add(taskName3); + PinotHelixTaskResourceManager spyMgr = Mockito.spy(mgr); + when(spyMgr.getTasks(taskType)).thenReturn(tasks); + + // Mock workflow-level components for getTaskStates + String helixJobQueueName = "TaskQueue_" + taskType; + WorkflowConfig workflowConfig = mock(WorkflowConfig.class); + when(taskDriver.getWorkflowConfig(helixJobQueueName)).thenReturn(workflowConfig); + + JobDag jobDag = mock(JobDag.class); + Set helixJobs = new HashSet<>(); + helixJobs.add(PinotHelixTaskResourceManager.getHelixJobName(taskName1)); + helixJobs.add(PinotHelixTaskResourceManager.getHelixJobName(taskName2)); + helixJobs.add(PinotHelixTaskResourceManager.getHelixJobName(taskName3)); + when(jobDag.getAllNodes()).thenReturn(helixJobs); + when(workflowConfig.getJobDag()).thenReturn(jobDag); + + WorkflowContext workflowContext = mock(WorkflowContext.class); + when(taskDriver.getWorkflowContext(helixJobQueueName)).thenReturn(workflowContext); + Map jobStatesMap = new HashMap<>(); + jobStatesMap.put(PinotHelixTaskResourceManager.getHelixJobName(taskName1), TaskState.IN_PROGRESS); // table1 + jobStatesMap.put(PinotHelixTaskResourceManager.getHelixJobName(taskName2), TaskState.COMPLETED); // table1 + jobStatesMap.put(PinotHelixTaskResourceManager.getHelixJobName(taskName3), TaskState.IN_PROGRESS); // table2 + when(workflowContext.getJobStates()).thenReturn(jobStatesMap); + + // Mock JobConfig and JobContext - different states for each task + mockTaskJobConfigAndContext(taskDriver, taskName1, TaskPartitionState.RUNNING); // table1 + mockTaskJobConfigAndContext(taskDriver, taskName2, TaskPartitionState.COMPLETED); // table1 + mockTaskJobConfigAndContext(taskDriver, taskName3, TaskPartitionState.RUNNING); // table2 + + // Mock subtask configs - taskName1 and taskName2 for table1, taskName3 for table2 + List subtaskConfigs1 = new ArrayList<>(); + PinotTaskConfig taskConfig1 = mock(PinotTaskConfig.class); + when(taskConfig1.getTableName()).thenReturn(table1); + subtaskConfigs1.add(taskConfig1); + when(spyMgr.getSubtaskConfigs(taskName1)).thenReturn(subtaskConfigs1); + + List subtaskConfigs2 = new ArrayList<>(); + PinotTaskConfig taskConfig2 = mock(PinotTaskConfig.class); + when(taskConfig2.getTableName()).thenReturn(table1); + subtaskConfigs2.add(taskConfig2); + when(spyMgr.getSubtaskConfigs(taskName2)).thenReturn(subtaskConfigs2); + + List subtaskConfigs3 = new ArrayList<>(); + PinotTaskConfig taskConfig3 = mock(PinotTaskConfig.class); + when(taskConfig3.getTableName()).thenReturn(table2); + subtaskConfigs3.add(taskConfig3); + when(spyMgr.getSubtaskConfigs(taskName3)).thenReturn(subtaskConfigs3); + + // Test filter by running state and table1 - should only return taskName1 + Map runningTable1Tasks = + spyMgr.getTaskCounts(taskType, "IN_PROGRESS", table1); + assertEquals(runningTable1Tasks.size(), 1); + assertTrue(runningTable1Tasks.containsKey(taskName1)); + assertFalse(runningTable1Tasks.containsKey(taskName2)); + assertFalse(runningTable1Tasks.containsKey(taskName3)); + + // Test filter by completed state and table1 - should only return taskName2 + Map completedTable1Tasks = + spyMgr.getTaskCounts(taskType, "COMPLETED", table1); + assertEquals(completedTable1Tasks.size(), 1); + assertFalse(completedTable1Tasks.containsKey(taskName1)); + assertTrue(completedTable1Tasks.containsKey(taskName2)); + assertFalse(completedTable1Tasks.containsKey(taskName3)); + + // Test filter by running state and table2 - should only return taskName3 + Map runningTable2Tasks = + spyMgr.getTaskCounts(taskType, "IN_PROGRESS", table2); + assertEquals(runningTable2Tasks.size(), 1); + assertFalse(runningTable2Tasks.containsKey(taskName1)); + assertFalse(runningTable2Tasks.containsKey(taskName2)); + assertTrue(runningTable2Tasks.containsKey(taskName3)); + + // Test filter by completed state and table2 - should return no tasks + Map completedTable2Tasks = + spyMgr.getTaskCounts(taskType, "COMPLETED", table2); + assertEquals(completedTable2Tasks.size(), 0); + + // Test filter by multiple states and table1 - should return taskName1 and taskName2 + Map multiStateTable1Tasks = + spyMgr.getTaskCounts(taskType, "IN_PROGRESS,COMPLETED", table1); + assertEquals(multiStateTable1Tasks.size(), 2); + assertTrue(multiStateTable1Tasks.containsKey(taskName1)); + assertTrue(multiStateTable1Tasks.containsKey(taskName2)); + assertFalse(multiStateTable1Tasks.containsKey(taskName3)); + } + + @Test + public void testGetTaskCountsWithTableFilterEdgeCases() { + String taskType = "TestTask"; + String taskName1 = "Task_TestTask_12345"; + + TaskDriver taskDriver = mock(TaskDriver.class); + PinotHelixResourceManager pinotHelixResourceManager = mock(PinotHelixResourceManager.class); + PinotHelixTaskResourceManager mgr = new PinotHelixTaskResourceManager(pinotHelixResourceManager, taskDriver); + + // Mock getTasks to return our test task name + Set tasks = new HashSet<>(); + tasks.add(taskName1); + PinotHelixTaskResourceManager spyMgr = Mockito.spy(mgr); + when(spyMgr.getTasks(taskType)).thenReturn(tasks); + + // Mock JobConfig and JobContext + mockTaskJobConfigAndContext(taskDriver, taskName1, TaskPartitionState.RUNNING); + + // Test with task that has null table name + List subtaskConfigsNullTable = new ArrayList<>(); + PinotTaskConfig taskConfigNullTable = mock(PinotTaskConfig.class); + when(taskConfigNullTable.getTableName()).thenReturn(null); + subtaskConfigsNullTable.add(taskConfigNullTable); + when(spyMgr.getSubtaskConfigs(taskName1)).thenReturn(subtaskConfigsNullTable); + + Map nullTableTasks = + spyMgr.getTaskCounts(taskType, null, "anyTable_OFFLINE"); + assertEquals(nullTableTasks.size(), 0); + + // Test with task that throws exception when getting subtask configs + when(spyMgr.getSubtaskConfigs(taskName1)).thenThrow(new RuntimeException("Test exception")); + + Map exceptionTasks = + spyMgr.getTaskCounts(taskType, null, "anyTable_OFFLINE"); + assertEquals(exceptionTasks.size(), 0); + + // Reset the mock to avoid exception in subsequent calls + Mockito.reset(spyMgr); + when(spyMgr.getTasks(taskType)).thenReturn(tasks); + mockTaskJobConfigAndContext(taskDriver, taskName1, TaskPartitionState.RUNNING); + + // Test with null state and null table (should return all tasks like original method) + when(spyMgr.getSubtaskConfigs(taskName1)).thenReturn(new ArrayList<>()); + Map allTasks = + spyMgr.getTaskCounts(taskType, null, null); + assertEquals(allTasks.size(), 1); + assertTrue(allTasks.containsKey(taskName1)); + } + + /** + * Helper method to mock JobConfig and JobContext for a task + */ + private void mockTaskJobConfigAndContext(TaskDriver taskDriver, String taskName, TaskPartitionState state) { + String helixJobName = PinotHelixTaskResourceManager.getHelixJobName(taskName); + JobConfig jobConfig = mock(JobConfig.class); + when(taskDriver.getJobConfig(helixJobName)).thenReturn(jobConfig); + Map taskConfigMap = new HashMap<>(); + taskConfigMap.put("taskId0", new TaskConfig("", new HashMap<>())); + when(jobConfig.getTaskConfigMap()).thenReturn(taskConfigMap); + JobContext jobContext = mock(JobContext.class); + when(taskDriver.getJobContext(helixJobName)).thenReturn(jobContext); + Map taskIdPartitionMap = new HashMap<>(); + taskIdPartitionMap.put("taskId0", 0); + when(jobContext.getTaskIdPartitionMap()).thenReturn(taskIdPartitionMap); + when(jobContext.getPartitionState(0)).thenReturn(state); + } } From 7ffd1e0cf25459745ed34377af19d843e8920f00 Mon Sep 17 00:00:00 2001 From: Yash Mayya Date: Tue, 29 Jul 2025 12:26:37 +0530 Subject: [PATCH 033/167] Move logical tables related classes from pinot-query-planner to pinot-core (#16443) --- .../pinot/broker/api/resources/PinotBrokerDebug.java | 2 +- .../BaseSingleStageBrokerRequestHandler.java | 10 +++++----- .../requesthandler/GrpcBrokerRequestHandler.java | 2 +- .../SingleConnectionBrokerRequestHandler.java | 2 +- .../pinot/broker/routing/BrokerRoutingManager.java | 6 +++--- .../routing/timeboundary/TimeBoundaryManager.java | 2 +- .../pinot/broker/broker/HelixBrokerStarterTest.java | 2 +- .../BaseSingleStageBrokerRequestHandlerTest.java | 2 +- .../timeboundary/TimeBoundaryManagerTest.java | 2 +- .../helix/core/retention/RetentionManager.java | 2 +- .../pinot/controller/util/BrokerServiceHelper.java | 2 +- .../controller/util/BrokerServiceHelperTest.java | 2 +- .../{transport => routing}/BaseTableRouteInfo.java | 2 +- .../routing}/ImplicitHybridTableRouteProvider.java | 8 ++------ .../pinot/core/routing}/LogicalTableRouteInfo.java | 9 +++------ .../core/routing}/LogicalTableRouteProvider.java | 10 ++++------ .../core/routing}/PhysicalTableRouteProvider.java | 6 +----- .../apache/pinot/core/routing/RoutingManager.java | 1 + .../core/{transport => routing}/TableRouteInfo.java | 7 ++++--- .../pinot/core/routing}/TableRouteProvider.java | 4 +--- .../timeboundary/MinTimeBoundaryStrategy.java | 3 +-- .../routing/{ => timeboundary}/TimeBoundaryInfo.java | 2 +- .../routing}/timeboundary/TimeBoundaryStrategy.java | 3 +-- .../timeboundary/TimeBoundaryStrategyService.java | 2 +- .../core/transport/ImplicitHybridTableRouteInfo.java | 3 ++- .../org/apache/pinot/core/transport/QueryRouter.java | 1 + .../pinot/core/routing}/BaseTableRouteTest.java | 12 ++++-------- ...itHybridTableRouteProviderCalculateRouteTest.java | 6 +----- ...citHybridTableRouteProviderGetTableRouteTest.java | 3 +-- .../LogicalTableRouteProviderCalculateRouteTest.java | 4 +--- .../LogicalTableRouteProviderGetRouteTest.java | 3 +-- .../core/routing}/MockRoutingManagerFactory.java | 9 ++------- .../timeboundary/MinTimeBoundaryStrategyTest.java | 3 +-- .../TimeBoundaryStrategyServiceTest.java | 2 +- ...ithTwoOfflineOneRealtimeTableIntegrationTest.java | 4 ++-- pinot-query-planner/pom.xml | 7 ++++++- .../planner/physical/DispatchablePlanFragment.java | 2 +- .../planner/physical/DispatchablePlanMetadata.java | 4 ++-- .../planner/physical/DispatchablePlanVisitor.java | 4 ++-- .../query/planner/physical/v2/TableScanMetadata.java | 2 +- .../v2/opt/rules/LeafStageWorkerAssignmentRule.java | 2 +- .../apache/pinot/query/routing/StageMetadata.java | 2 +- .../apache/pinot/query/routing/WorkerManager.java | 8 ++++---- .../apache/pinot/query/QueryEnvironmentTestBase.java | 2 +- .../runtime/plan/server/ServerPlanRequestUtils.java | 2 +- .../pinot/query/service/server/QueryServerTest.java | 2 +- .../tsdb/planner/TimeSeriesQueryEnvironment.java | 2 +- .../tsdb/planner/physical/TableScanVisitor.java | 4 ++-- 48 files changed, 80 insertions(+), 106 deletions(-) rename pinot-core/src/main/java/org/apache/pinot/core/{transport => routing}/BaseTableRouteInfo.java (97%) rename {pinot-query-planner/src/main/java/org/apache/pinot/query/routing/table => pinot-core/src/main/java/org/apache/pinot/core/routing}/ImplicitHybridTableRouteProvider.java (96%) rename {pinot-query-planner/src/main/java/org/apache/pinot/query/routing/table => pinot-core/src/main/java/org/apache/pinot/core/routing}/LogicalTableRouteInfo.java (97%) rename {pinot-query-planner/src/main/java/org/apache/pinot/query/routing/table => pinot-core/src/main/java/org/apache/pinot/core/routing}/LogicalTableRouteProvider.java (95%) rename {pinot-query-planner/src/main/java/org/apache/pinot/query/routing/table => pinot-core/src/main/java/org/apache/pinot/core/routing}/PhysicalTableRouteProvider.java (94%) rename pinot-core/src/main/java/org/apache/pinot/core/{transport => routing}/TableRouteInfo.java (96%) rename {pinot-query-planner/src/main/java/org/apache/pinot/query/routing/table => pinot-core/src/main/java/org/apache/pinot/core/routing}/TableRouteProvider.java (93%) rename {pinot-query-planner/src/main/java/org/apache/pinot/query => pinot-core/src/main/java/org/apache/pinot/core/routing}/timeboundary/MinTimeBoundaryStrategy.java (97%) rename pinot-core/src/main/java/org/apache/pinot/core/routing/{ => timeboundary}/TimeBoundaryInfo.java (96%) rename {pinot-query-planner/src/main/java/org/apache/pinot/query => pinot-core/src/main/java/org/apache/pinot/core/routing}/timeboundary/TimeBoundaryStrategy.java (95%) rename {pinot-query-planner/src/main/java/org/apache/pinot/query => pinot-core/src/main/java/org/apache/pinot/core/routing}/timeboundary/TimeBoundaryStrategyService.java (97%) rename {pinot-query-planner/src/test/java/org/apache/pinot/query/routing/table => pinot-core/src/test/java/org/apache/pinot/core/routing}/BaseTableRouteTest.java (97%) rename {pinot-query-planner/src/test/java/org/apache/pinot/query/routing/table => pinot-core/src/test/java/org/apache/pinot/core/routing}/ImplicitHybridTableRouteProviderCalculateRouteTest.java (98%) rename {pinot-query-planner/src/test/java/org/apache/pinot/query/routing/table => pinot-core/src/test/java/org/apache/pinot/core/routing}/ImplicitHybridTableRouteProviderGetTableRouteTest.java (99%) rename {pinot-query-planner/src/test/java/org/apache/pinot/query/routing/table => pinot-core/src/test/java/org/apache/pinot/core/routing}/LogicalTableRouteProviderCalculateRouteTest.java (97%) rename {pinot-query-planner/src/test/java/org/apache/pinot/query/routing/table => pinot-core/src/test/java/org/apache/pinot/core/routing}/LogicalTableRouteProviderGetRouteTest.java (99%) rename {pinot-query-planner/src/test/java/org/apache/pinot/query/testutils => pinot-core/src/test/java/org/apache/pinot/core/routing}/MockRoutingManagerFactory.java (96%) rename {pinot-query-planner/src/test/java/org/apache/pinot/query => pinot-core/src/test/java/org/apache/pinot/core/routing}/timeboundary/MinTimeBoundaryStrategyTest.java (98%) rename {pinot-query-planner/src/test/java/org/apache/pinot/query => pinot-core/src/test/java/org/apache/pinot/core/routing}/timeboundary/TimeBoundaryStrategyServiceTest.java (96%) diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotBrokerDebug.java b/pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotBrokerDebug.java index 0d534476d7cc..19733099b505 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotBrokerDebug.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotBrokerDebug.java @@ -56,7 +56,7 @@ import org.apache.pinot.core.auth.TargetType; import org.apache.pinot.core.routing.RoutingTable; import org.apache.pinot.core.routing.ServerRouteInfo; -import org.apache.pinot.core.routing.TimeBoundaryInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; import org.apache.pinot.core.transport.ServerInstance; import org.apache.pinot.core.transport.server.routing.stats.ServerRoutingStatsEntry; import org.apache.pinot.core.transport.server.routing.stats.ServerRoutingStatsManager; diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java index 8ef923f63bbe..4f409e8ad302 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java @@ -90,14 +90,14 @@ import org.apache.pinot.core.query.reduce.GapfillProcessorFactory; import org.apache.pinot.core.query.request.context.QueryContext; import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils; -import org.apache.pinot.core.routing.TimeBoundaryInfo; +import org.apache.pinot.core.routing.ImplicitHybridTableRouteProvider; +import org.apache.pinot.core.routing.LogicalTableRouteProvider; +import org.apache.pinot.core.routing.TableRouteInfo; +import org.apache.pinot.core.routing.TableRouteProvider; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; import org.apache.pinot.core.transport.ServerInstance; -import org.apache.pinot.core.transport.TableRouteInfo; import org.apache.pinot.core.util.GapfillUtils; import org.apache.pinot.query.parser.utils.ParserUtils; -import org.apache.pinot.query.routing.table.ImplicitHybridTableRouteProvider; -import org.apache.pinot.query.routing.table.LogicalTableRouteProvider; -import org.apache.pinot.query.routing.table.TableRouteProvider; import org.apache.pinot.segment.local.function.GroovyFunctionEvaluator; import org.apache.pinot.spi.accounting.ThreadExecutionContext; import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant; diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java index 0c832085ee26..5b9649e92cba 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java @@ -39,9 +39,9 @@ import org.apache.pinot.common.utils.grpc.ServerGrpcRequestBuilder; import org.apache.pinot.core.query.reduce.StreamingReduceService; import org.apache.pinot.core.routing.ServerRouteInfo; +import org.apache.pinot.core.routing.TableRouteInfo; import org.apache.pinot.core.transport.ServerInstance; import org.apache.pinot.core.transport.ServerRoutingInstance; -import org.apache.pinot.core.transport.TableRouteInfo; import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.env.PinotConfiguration; diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java index e636a2ce5b7d..ebbab31b3bee 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java @@ -39,13 +39,13 @@ import org.apache.pinot.common.response.broker.QueryProcessingException; import org.apache.pinot.common.utils.config.QueryOptionsUtils; import org.apache.pinot.core.query.reduce.BrokerReduceService; +import org.apache.pinot.core.routing.TableRouteInfo; import org.apache.pinot.core.transport.AsyncQueryResponse; import org.apache.pinot.core.transport.QueryResponse; import org.apache.pinot.core.transport.QueryRouter; import org.apache.pinot.core.transport.ServerInstance; import org.apache.pinot.core.transport.ServerResponse; import org.apache.pinot.core.transport.ServerRoutingInstance; -import org.apache.pinot.core.transport.TableRouteInfo; import org.apache.pinot.core.transport.server.routing.stats.ServerRoutingStatsManager; import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant; import org.apache.pinot.spi.env.PinotConfiguration; diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/BrokerRoutingManager.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/BrokerRoutingManager.java index b2edab2e20f5..d5d345d1abd4 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/BrokerRoutingManager.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/BrokerRoutingManager.java @@ -64,11 +64,11 @@ import org.apache.pinot.core.routing.ServerRouteInfo; import org.apache.pinot.core.routing.TablePartitionInfo; import org.apache.pinot.core.routing.TablePartitionReplicatedServersInfo; -import org.apache.pinot.core.routing.TimeBoundaryInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryStrategy; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryStrategyService; import org.apache.pinot.core.transport.ServerInstance; import org.apache.pinot.core.transport.server.routing.stats.ServerRoutingStatsManager; -import org.apache.pinot.query.timeboundary.TimeBoundaryStrategy; -import org.apache.pinot.query.timeboundary.TimeBoundaryStrategyService; import org.apache.pinot.spi.config.table.ColumnPartitionConfig; import org.apache.pinot.spi.config.table.QueryConfig; import org.apache.pinot.spi.config.table.SegmentPartitionConfig; diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/timeboundary/TimeBoundaryManager.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/timeboundary/TimeBoundaryManager.java index 13b2ead1e41f..b7987a83fd71 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/timeboundary/TimeBoundaryManager.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/timeboundary/TimeBoundaryManager.java @@ -36,7 +36,7 @@ import org.apache.pinot.common.metadata.segment.SegmentZKMetadata; import org.apache.pinot.common.metrics.BrokerGauge; import org.apache.pinot.common.metrics.BrokerMetrics; -import org.apache.pinot.core.routing.TimeBoundaryInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.data.DateTimeFieldSpec; diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/broker/HelixBrokerStarterTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/broker/HelixBrokerStarterTest.java index 7a1f1d7c7f4b..3310a0e8959b 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/broker/HelixBrokerStarterTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/broker/HelixBrokerStarterTest.java @@ -34,7 +34,7 @@ import org.apache.pinot.controller.helix.ControllerTest; import org.apache.pinot.controller.utils.SegmentMetadataMockUtils; import org.apache.pinot.core.routing.RoutingTable; -import org.apache.pinot.core.routing.TimeBoundaryInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.data.FieldSpec; diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandlerTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandlerTest.java index a3e80e0e882e..9c315715531f 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandlerTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandlerTest.java @@ -36,8 +36,8 @@ import org.apache.pinot.common.response.broker.BrokerResponseNative; import org.apache.pinot.core.routing.RoutingTable; import org.apache.pinot.core.routing.ServerRouteInfo; +import org.apache.pinot.core.routing.TableRouteInfo; import org.apache.pinot.core.transport.ServerInstance; -import org.apache.pinot.core.transport.TableRouteInfo; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TenantConfig; import org.apache.pinot.spi.env.PinotConfiguration; diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/timeboundary/TimeBoundaryManagerTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/timeboundary/TimeBoundaryManagerTest.java index 96a3472d5e5b..bed9f899d3b9 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/timeboundary/TimeBoundaryManagerTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/timeboundary/TimeBoundaryManagerTest.java @@ -34,7 +34,7 @@ import org.apache.pinot.common.metadata.segment.SegmentZKMetadata; import org.apache.pinot.common.metrics.BrokerMetrics; import org.apache.pinot.controller.helix.ControllerTest; -import org.apache.pinot.core.routing.TimeBoundaryInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.data.FieldSpec; diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/retention/RetentionManager.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/retention/RetentionManager.java index 984bfbd52726..d647574eee41 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/retention/RetentionManager.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/retention/RetentionManager.java @@ -50,7 +50,7 @@ import org.apache.pinot.controller.helix.core.retention.strategy.RetentionStrategy; import org.apache.pinot.controller.helix.core.retention.strategy.TimeRetentionStrategy; import org.apache.pinot.controller.util.BrokerServiceHelper; -import org.apache.pinot.core.routing.TimeBoundaryInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; import org.apache.pinot.spi.config.table.SegmentsValidationAndRetentionConfig; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/util/BrokerServiceHelper.java b/pinot-controller/src/main/java/org/apache/pinot/controller/util/BrokerServiceHelper.java index 18ae860aaca3..25317492a12d 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/util/BrokerServiceHelper.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/util/BrokerServiceHelper.java @@ -31,7 +31,7 @@ import org.apache.pinot.common.auth.AuthProviderUtils; import org.apache.pinot.controller.ControllerConf; import org.apache.pinot.controller.helix.core.PinotHelixResourceManager; -import org.apache.pinot.core.routing.TimeBoundaryInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; import org.apache.pinot.spi.auth.AuthProvider; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.utils.CommonConstants; diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/util/BrokerServiceHelperTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/util/BrokerServiceHelperTest.java index 2b413f203819..428c4e4491fe 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/util/BrokerServiceHelperTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/util/BrokerServiceHelperTest.java @@ -25,7 +25,7 @@ import org.apache.helix.model.InstanceConfig; import org.apache.pinot.controller.ControllerConf; import org.apache.pinot.controller.helix.core.PinotHelixResourceManager; -import org.apache.pinot.core.routing.TimeBoundaryInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.utils.builder.TableConfigBuilder; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/transport/BaseTableRouteInfo.java b/pinot-core/src/main/java/org/apache/pinot/core/routing/BaseTableRouteInfo.java similarity index 97% rename from pinot-core/src/main/java/org/apache/pinot/core/transport/BaseTableRouteInfo.java rename to pinot-core/src/main/java/org/apache/pinot/core/routing/BaseTableRouteInfo.java index ff55ba15863d..43b244d08acf 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/transport/BaseTableRouteInfo.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/routing/BaseTableRouteInfo.java @@ -16,7 +16,7 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.pinot.core.transport; +package org.apache.pinot.core.routing; public abstract class BaseTableRouteInfo implements TableRouteInfo { diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/table/ImplicitHybridTableRouteProvider.java b/pinot-core/src/main/java/org/apache/pinot/core/routing/ImplicitHybridTableRouteProvider.java similarity index 96% rename from pinot-query-planner/src/main/java/org/apache/pinot/query/routing/table/ImplicitHybridTableRouteProvider.java rename to pinot-core/src/main/java/org/apache/pinot/core/routing/ImplicitHybridTableRouteProvider.java index 48553c328d81..14c8e726f81d 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/table/ImplicitHybridTableRouteProvider.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/routing/ImplicitHybridTableRouteProvider.java @@ -16,7 +16,7 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.pinot.query.routing.table; +package org.apache.pinot.core.routing; import com.google.common.base.Preconditions; import java.util.ArrayList; @@ -24,13 +24,9 @@ import java.util.Map; import org.apache.pinot.common.config.provider.TableCache; import org.apache.pinot.common.request.BrokerRequest; -import org.apache.pinot.core.routing.RoutingManager; -import org.apache.pinot.core.routing.RoutingTable; -import org.apache.pinot.core.routing.ServerRouteInfo; -import org.apache.pinot.core.routing.TimeBoundaryInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; import org.apache.pinot.core.transport.ImplicitHybridTableRouteInfo; import org.apache.pinot.core.transport.ServerInstance; -import org.apache.pinot.core.transport.TableRouteInfo; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.utils.builder.TableNameBuilder; import org.slf4j.Logger; diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/table/LogicalTableRouteInfo.java b/pinot-core/src/main/java/org/apache/pinot/core/routing/LogicalTableRouteInfo.java similarity index 97% rename from pinot-query-planner/src/main/java/org/apache/pinot/query/routing/table/LogicalTableRouteInfo.java rename to pinot-core/src/main/java/org/apache/pinot/core/routing/LogicalTableRouteInfo.java index 612f8e79cd3e..39a3a709c28f 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/table/LogicalTableRouteInfo.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/routing/LogicalTableRouteInfo.java @@ -16,7 +16,7 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.pinot.query.routing.table; +package org.apache.pinot.core.routing; import java.util.ArrayList; import java.util.HashMap; @@ -29,14 +29,11 @@ import org.apache.pinot.common.request.BrokerRequest; import org.apache.pinot.common.request.InstanceRequest; import org.apache.pinot.common.request.TableSegmentsInfo; -import org.apache.pinot.core.routing.ServerRouteInfo; -import org.apache.pinot.core.routing.TimeBoundaryInfo; -import org.apache.pinot.core.transport.BaseTableRouteInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryStrategy; import org.apache.pinot.core.transport.ImplicitHybridTableRouteInfo; import org.apache.pinot.core.transport.ServerInstance; import org.apache.pinot.core.transport.ServerRoutingInstance; -import org.apache.pinot.core.transport.TableRouteInfo; -import org.apache.pinot.query.timeboundary.TimeBoundaryStrategy; import org.apache.pinot.spi.config.table.QueryConfig; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/table/LogicalTableRouteProvider.java b/pinot-core/src/main/java/org/apache/pinot/core/routing/LogicalTableRouteProvider.java similarity index 95% rename from pinot-query-planner/src/main/java/org/apache/pinot/query/routing/table/LogicalTableRouteProvider.java rename to pinot-core/src/main/java/org/apache/pinot/core/routing/LogicalTableRouteProvider.java index bdce62a93838..cf33604c382b 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/table/LogicalTableRouteProvider.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/routing/LogicalTableRouteProvider.java @@ -16,19 +16,17 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.pinot.query.routing.table; +package org.apache.pinot.core.routing; import com.google.common.base.Preconditions; import java.util.ArrayList; import java.util.List; import org.apache.pinot.common.config.provider.TableCache; import org.apache.pinot.common.request.BrokerRequest; -import org.apache.pinot.core.routing.RoutingManager; -import org.apache.pinot.core.routing.TimeBoundaryInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryStrategy; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryStrategyService; import org.apache.pinot.core.transport.ImplicitHybridTableRouteInfo; -import org.apache.pinot.core.transport.TableRouteInfo; -import org.apache.pinot.query.timeboundary.TimeBoundaryStrategy; -import org.apache.pinot.query.timeboundary.TimeBoundaryStrategyService; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.data.LogicalTableConfig; diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/table/PhysicalTableRouteProvider.java b/pinot-core/src/main/java/org/apache/pinot/core/routing/PhysicalTableRouteProvider.java similarity index 94% rename from pinot-query-planner/src/main/java/org/apache/pinot/query/routing/table/PhysicalTableRouteProvider.java rename to pinot-core/src/main/java/org/apache/pinot/core/routing/PhysicalTableRouteProvider.java index 04d09fd63473..b175f97f5a7a 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/table/PhysicalTableRouteProvider.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/routing/PhysicalTableRouteProvider.java @@ -16,19 +16,15 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.pinot.query.routing.table; +package org.apache.pinot.core.routing; import com.google.common.base.Preconditions; import java.util.ArrayList; import java.util.List; import java.util.Map; import org.apache.pinot.common.request.BrokerRequest; -import org.apache.pinot.core.routing.RoutingManager; -import org.apache.pinot.core.routing.RoutingTable; -import org.apache.pinot.core.routing.ServerRouteInfo; import org.apache.pinot.core.transport.ImplicitHybridTableRouteInfo; import org.apache.pinot.core.transport.ServerInstance; -import org.apache.pinot.core.transport.TableRouteInfo; /** diff --git a/pinot-core/src/main/java/org/apache/pinot/core/routing/RoutingManager.java b/pinot-core/src/main/java/org/apache/pinot/core/routing/RoutingManager.java index f3fad641dab9..e0e00cdd04e6 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/routing/RoutingManager.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/routing/RoutingManager.java @@ -23,6 +23,7 @@ import java.util.Set; import javax.annotation.Nullable; import org.apache.pinot.common.request.BrokerRequest; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; import org.apache.pinot.core.transport.ServerInstance; import org.apache.pinot.spi.annotations.InterfaceAudience; import org.apache.pinot.spi.annotations.InterfaceStability; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/transport/TableRouteInfo.java b/pinot-core/src/main/java/org/apache/pinot/core/routing/TableRouteInfo.java similarity index 96% rename from pinot-core/src/main/java/org/apache/pinot/core/transport/TableRouteInfo.java rename to pinot-core/src/main/java/org/apache/pinot/core/routing/TableRouteInfo.java index b969d20e5234..cf8220cbb2a2 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/transport/TableRouteInfo.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/routing/TableRouteInfo.java @@ -16,7 +16,7 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.pinot.core.transport; +package org.apache.pinot.core.routing; import java.util.List; import java.util.Map; @@ -24,8 +24,9 @@ import javax.annotation.Nullable; import org.apache.pinot.common.request.BrokerRequest; import org.apache.pinot.common.request.InstanceRequest; -import org.apache.pinot.core.routing.ServerRouteInfo; -import org.apache.pinot.core.routing.TimeBoundaryInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; +import org.apache.pinot.core.transport.ServerInstance; +import org.apache.pinot.core.transport.ServerRoutingInstance; import org.apache.pinot.spi.config.table.QueryConfig; import org.apache.pinot.spi.config.table.TableConfig; diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/table/TableRouteProvider.java b/pinot-core/src/main/java/org/apache/pinot/core/routing/TableRouteProvider.java similarity index 93% rename from pinot-query-planner/src/main/java/org/apache/pinot/query/routing/table/TableRouteProvider.java rename to pinot-core/src/main/java/org/apache/pinot/core/routing/TableRouteProvider.java index d4707f9975ad..22912b7f4ddd 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/table/TableRouteProvider.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/routing/TableRouteProvider.java @@ -16,12 +16,10 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.pinot.query.routing.table; +package org.apache.pinot.core.routing; import org.apache.pinot.common.config.provider.TableCache; import org.apache.pinot.common.request.BrokerRequest; -import org.apache.pinot.core.routing.RoutingManager; -import org.apache.pinot.core.transport.TableRouteInfo; /** diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/timeboundary/MinTimeBoundaryStrategy.java b/pinot-core/src/main/java/org/apache/pinot/core/routing/timeboundary/MinTimeBoundaryStrategy.java similarity index 97% rename from pinot-query-planner/src/main/java/org/apache/pinot/query/timeboundary/MinTimeBoundaryStrategy.java rename to pinot-core/src/main/java/org/apache/pinot/core/routing/timeboundary/MinTimeBoundaryStrategy.java index e204c917daf0..83de6b2fc749 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/timeboundary/MinTimeBoundaryStrategy.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/routing/timeboundary/MinTimeBoundaryStrategy.java @@ -16,7 +16,7 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.pinot.query.timeboundary; +package org.apache.pinot.core.routing.timeboundary; import com.google.auto.service.AutoService; import com.google.common.base.Preconditions; @@ -25,7 +25,6 @@ import java.util.Map; import org.apache.pinot.common.config.provider.TableCache; import org.apache.pinot.core.routing.RoutingManager; -import org.apache.pinot.core.routing.TimeBoundaryInfo; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.data.DateTimeFieldSpec; import org.apache.pinot.spi.data.DateTimeFormatSpec; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/routing/TimeBoundaryInfo.java b/pinot-core/src/main/java/org/apache/pinot/core/routing/timeboundary/TimeBoundaryInfo.java similarity index 96% rename from pinot-core/src/main/java/org/apache/pinot/core/routing/TimeBoundaryInfo.java rename to pinot-core/src/main/java/org/apache/pinot/core/routing/timeboundary/TimeBoundaryInfo.java index 07ae45335595..f9382c5d5dde 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/routing/TimeBoundaryInfo.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/routing/timeboundary/TimeBoundaryInfo.java @@ -16,7 +16,7 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.pinot.core.routing; +package org.apache.pinot.core.routing.timeboundary; import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonProperty; diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/timeboundary/TimeBoundaryStrategy.java b/pinot-core/src/main/java/org/apache/pinot/core/routing/timeboundary/TimeBoundaryStrategy.java similarity index 95% rename from pinot-query-planner/src/main/java/org/apache/pinot/query/timeboundary/TimeBoundaryStrategy.java rename to pinot-core/src/main/java/org/apache/pinot/core/routing/timeboundary/TimeBoundaryStrategy.java index f2215b29b226..442c0dfd2168 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/timeboundary/TimeBoundaryStrategy.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/routing/timeboundary/TimeBoundaryStrategy.java @@ -16,12 +16,11 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.pinot.query.timeboundary; +package org.apache.pinot.core.routing.timeboundary; import java.util.List; import org.apache.pinot.common.config.provider.TableCache; import org.apache.pinot.core.routing.RoutingManager; -import org.apache.pinot.core.routing.TimeBoundaryInfo; import org.apache.pinot.spi.data.LogicalTableConfig; diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/timeboundary/TimeBoundaryStrategyService.java b/pinot-core/src/main/java/org/apache/pinot/core/routing/timeboundary/TimeBoundaryStrategyService.java similarity index 97% rename from pinot-query-planner/src/main/java/org/apache/pinot/query/timeboundary/TimeBoundaryStrategyService.java rename to pinot-core/src/main/java/org/apache/pinot/core/routing/timeboundary/TimeBoundaryStrategyService.java index 0e3f6af4aa0f..51b952df84ff 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/timeboundary/TimeBoundaryStrategyService.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/routing/timeboundary/TimeBoundaryStrategyService.java @@ -16,7 +16,7 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.pinot.query.timeboundary; +package org.apache.pinot.core.routing.timeboundary; import java.util.Map; import java.util.ServiceLoader; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/transport/ImplicitHybridTableRouteInfo.java b/pinot-core/src/main/java/org/apache/pinot/core/transport/ImplicitHybridTableRouteInfo.java index 00199c272698..90cc9089a2b8 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/transport/ImplicitHybridTableRouteInfo.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/transport/ImplicitHybridTableRouteInfo.java @@ -26,8 +26,9 @@ import org.apache.commons.collections4.CollectionUtils; import org.apache.pinot.common.request.BrokerRequest; import org.apache.pinot.common.request.InstanceRequest; +import org.apache.pinot.core.routing.BaseTableRouteInfo; import org.apache.pinot.core.routing.ServerRouteInfo; -import org.apache.pinot.core.routing.TimeBoundaryInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; import org.apache.pinot.spi.config.table.QueryConfig; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/transport/QueryRouter.java b/pinot-core/src/main/java/org/apache/pinot/core/transport/QueryRouter.java index 655a5796e233..a01f248751ef 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/transport/QueryRouter.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/transport/QueryRouter.java @@ -33,6 +33,7 @@ import org.apache.pinot.common.request.InstanceRequest; import org.apache.pinot.common.utils.config.QueryOptionsUtils; import org.apache.pinot.core.routing.ServerRouteInfo; +import org.apache.pinot.core.routing.TableRouteInfo; import org.apache.pinot.core.transport.server.routing.stats.ServerRoutingStatsManager; import org.apache.pinot.spi.config.table.TableType; import org.slf4j.Logger; diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/table/BaseTableRouteTest.java b/pinot-core/src/test/java/org/apache/pinot/core/routing/BaseTableRouteTest.java similarity index 97% rename from pinot-query-planner/src/test/java/org/apache/pinot/query/routing/table/BaseTableRouteTest.java rename to pinot-core/src/test/java/org/apache/pinot/core/routing/BaseTableRouteTest.java index 71901f843215..c3d70431a74d 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/table/BaseTableRouteTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/routing/BaseTableRouteTest.java @@ -16,7 +16,7 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.pinot.query.routing.table; +package org.apache.pinot.core.routing; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; @@ -29,14 +29,10 @@ import org.apache.pinot.common.config.provider.TableCache; import org.apache.pinot.common.request.BrokerRequest; import org.apache.pinot.common.request.PinotQuery; -import org.apache.pinot.core.routing.RoutingManager; -import org.apache.pinot.core.routing.ServerRouteInfo; -import org.apache.pinot.core.routing.TimeBoundaryInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryStrategy; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryStrategyService; import org.apache.pinot.core.transport.ServerInstance; -import org.apache.pinot.core.transport.TableRouteInfo; -import org.apache.pinot.query.testutils.MockRoutingManagerFactory; -import org.apache.pinot.query.timeboundary.TimeBoundaryStrategy; -import org.apache.pinot.query.timeboundary.TimeBoundaryStrategyService; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.data.LogicalTableConfig; diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/table/ImplicitHybridTableRouteProviderCalculateRouteTest.java b/pinot-core/src/test/java/org/apache/pinot/core/routing/ImplicitHybridTableRouteProviderCalculateRouteTest.java similarity index 98% rename from pinot-query-planner/src/test/java/org/apache/pinot/query/routing/table/ImplicitHybridTableRouteProviderCalculateRouteTest.java rename to pinot-core/src/test/java/org/apache/pinot/core/routing/ImplicitHybridTableRouteProviderCalculateRouteTest.java index be8606d06890..e9cb15fbba09 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/table/ImplicitHybridTableRouteProviderCalculateRouteTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/routing/ImplicitHybridTableRouteProviderCalculateRouteTest.java @@ -16,7 +16,7 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.pinot.query.routing.table; +package org.apache.pinot.core.routing; import java.util.ArrayList; import java.util.List; @@ -24,12 +24,8 @@ import java.util.Set; import org.apache.pinot.common.request.BrokerRequest; import org.apache.pinot.common.request.InstanceRequest; -import org.apache.pinot.core.routing.RoutingManager; -import org.apache.pinot.core.routing.RoutingTable; -import org.apache.pinot.core.routing.ServerRouteInfo; import org.apache.pinot.core.transport.ServerInstance; import org.apache.pinot.core.transport.ServerRoutingInstance; -import org.apache.pinot.core.transport.TableRouteInfo; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.query.QueryThreadContext; import org.apache.pinot.spi.utils.builder.TableNameBuilder; diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/table/ImplicitHybridTableRouteProviderGetTableRouteTest.java b/pinot-core/src/test/java/org/apache/pinot/core/routing/ImplicitHybridTableRouteProviderGetTableRouteTest.java similarity index 99% rename from pinot-query-planner/src/test/java/org/apache/pinot/query/routing/table/ImplicitHybridTableRouteProviderGetTableRouteTest.java rename to pinot-core/src/test/java/org/apache/pinot/core/routing/ImplicitHybridTableRouteProviderGetTableRouteTest.java index ddd34ec9da3d..4f12019b8387 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/table/ImplicitHybridTableRouteProviderGetTableRouteTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/routing/ImplicitHybridTableRouteProviderGetTableRouteTest.java @@ -16,12 +16,11 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.pinot.query.routing.table; +package org.apache.pinot.core.routing; import org.apache.pinot.common.response.BrokerResponse; import org.apache.pinot.common.response.broker.BrokerResponseNative; import org.apache.pinot.common.response.broker.QueryProcessingException; -import org.apache.pinot.core.transport.TableRouteInfo; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.exception.QueryErrorCode; diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/table/LogicalTableRouteProviderCalculateRouteTest.java b/pinot-core/src/test/java/org/apache/pinot/core/routing/LogicalTableRouteProviderCalculateRouteTest.java similarity index 97% rename from pinot-query-planner/src/test/java/org/apache/pinot/query/routing/table/LogicalTableRouteProviderCalculateRouteTest.java rename to pinot-core/src/test/java/org/apache/pinot/core/routing/LogicalTableRouteProviderCalculateRouteTest.java index f76a3bfdb972..70311857fb1b 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/table/LogicalTableRouteProviderCalculateRouteTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/routing/LogicalTableRouteProviderCalculateRouteTest.java @@ -16,15 +16,13 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.pinot.query.routing.table; +package org.apache.pinot.core.routing; import java.util.Map; import java.util.Set; import org.apache.pinot.common.request.InstanceRequest; -import org.apache.pinot.core.routing.ServerRouteInfo; import org.apache.pinot.core.transport.ServerInstance; import org.apache.pinot.core.transport.ServerRoutingInstance; -import org.apache.pinot.core.transport.TableRouteInfo; import org.apache.pinot.spi.query.QueryThreadContext; import org.testng.annotations.AfterMethod; import org.testng.annotations.BeforeMethod; diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/table/LogicalTableRouteProviderGetRouteTest.java b/pinot-core/src/test/java/org/apache/pinot/core/routing/LogicalTableRouteProviderGetRouteTest.java similarity index 99% rename from pinot-query-planner/src/test/java/org/apache/pinot/query/routing/table/LogicalTableRouteProviderGetRouteTest.java rename to pinot-core/src/test/java/org/apache/pinot/core/routing/LogicalTableRouteProviderGetRouteTest.java index d9d0a9fa5152..b846ef1bb0f4 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/routing/table/LogicalTableRouteProviderGetRouteTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/routing/LogicalTableRouteProviderGetRouteTest.java @@ -16,14 +16,13 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.pinot.query.routing.table; +package org.apache.pinot.core.routing; import com.google.common.collect.ImmutableList; import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; -import org.apache.pinot.core.transport.TableRouteInfo; import org.apache.pinot.spi.data.LogicalTableConfig; import org.apache.pinot.spi.data.PhysicalTableConfig; import org.apache.pinot.spi.data.TimeBoundaryConfig; diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/testutils/MockRoutingManagerFactory.java b/pinot-core/src/test/java/org/apache/pinot/core/routing/MockRoutingManagerFactory.java similarity index 96% rename from pinot-query-planner/src/test/java/org/apache/pinot/query/testutils/MockRoutingManagerFactory.java rename to pinot-core/src/test/java/org/apache/pinot/core/routing/MockRoutingManagerFactory.java index bf6bca907703..fba9f7945075 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/testutils/MockRoutingManagerFactory.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/routing/MockRoutingManagerFactory.java @@ -16,7 +16,7 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.pinot.query.testutils; +package org.apache.pinot.core.routing; import com.google.common.collect.Maps; import java.util.ArrayList; @@ -31,12 +31,7 @@ import org.apache.helix.zookeeper.datamodel.ZNRecord; import org.apache.pinot.common.config.provider.TableCache; import org.apache.pinot.common.request.BrokerRequest; -import org.apache.pinot.core.routing.RoutingManager; -import org.apache.pinot.core.routing.RoutingTable; -import org.apache.pinot.core.routing.ServerRouteInfo; -import org.apache.pinot.core.routing.TablePartitionInfo; -import org.apache.pinot.core.routing.TablePartitionReplicatedServersInfo; -import org.apache.pinot.core.routing.TimeBoundaryInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; import org.apache.pinot.core.transport.ServerInstance; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.data.Schema; diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/timeboundary/MinTimeBoundaryStrategyTest.java b/pinot-core/src/test/java/org/apache/pinot/core/routing/timeboundary/MinTimeBoundaryStrategyTest.java similarity index 98% rename from pinot-query-planner/src/test/java/org/apache/pinot/query/timeboundary/MinTimeBoundaryStrategyTest.java rename to pinot-core/src/test/java/org/apache/pinot/core/routing/timeboundary/MinTimeBoundaryStrategyTest.java index 7e0ce2d2feab..36e6b58c6dbd 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/timeboundary/MinTimeBoundaryStrategyTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/routing/timeboundary/MinTimeBoundaryStrategyTest.java @@ -16,13 +16,12 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.pinot.query.timeboundary; +package org.apache.pinot.core.routing.timeboundary; import java.util.List; import java.util.Map; import org.apache.pinot.common.config.provider.TableCache; import org.apache.pinot.core.routing.RoutingManager; -import org.apache.pinot.core.routing.TimeBoundaryInfo; import org.apache.pinot.spi.config.table.SegmentsValidationAndRetentionConfig; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.data.DateTimeFieldSpec; diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/timeboundary/TimeBoundaryStrategyServiceTest.java b/pinot-core/src/test/java/org/apache/pinot/core/routing/timeboundary/TimeBoundaryStrategyServiceTest.java similarity index 96% rename from pinot-query-planner/src/test/java/org/apache/pinot/query/timeboundary/TimeBoundaryStrategyServiceTest.java rename to pinot-core/src/test/java/org/apache/pinot/core/routing/timeboundary/TimeBoundaryStrategyServiceTest.java index 05993a437f44..0d72e0502674 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/timeboundary/TimeBoundaryStrategyServiceTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/routing/timeboundary/TimeBoundaryStrategyServiceTest.java @@ -16,7 +16,7 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.pinot.query.timeboundary; +package org.apache.pinot.core.routing.timeboundary; import org.testng.annotations.Test; diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/logicaltable/LogicalTableWithTwoOfflineOneRealtimeTableIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/logicaltable/LogicalTableWithTwoOfflineOneRealtimeTableIntegrationTest.java index cafa33a56d7d..a11f248b130c 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/logicaltable/LogicalTableWithTwoOfflineOneRealtimeTableIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/logicaltable/LogicalTableWithTwoOfflineOneRealtimeTableIntegrationTest.java @@ -22,8 +22,8 @@ import java.io.IOException; import java.util.List; import java.util.Map; -import org.apache.pinot.query.timeboundary.TimeBoundaryStrategy; -import org.apache.pinot.query.timeboundary.TimeBoundaryStrategyService; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryStrategy; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryStrategyService; import org.apache.pinot.spi.data.LogicalTableConfig; import org.apache.pinot.spi.utils.builder.TableNameBuilder; import org.testng.annotations.Test; diff --git a/pinot-query-planner/pom.xml b/pinot-query-planner/pom.xml index 936213bda01e..7958fd65f376 100644 --- a/pinot-query-planner/pom.xml +++ b/pinot-query-planner/pom.xml @@ -40,7 +40,6 @@ org.apache.pinot pinot-core - org.codehaus.janino janino @@ -54,6 +53,12 @@ value-annotations + + org.apache.pinot + pinot-core + test + test-jar + org.testng testng diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanFragment.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanFragment.java index 61b50ca8ebd6..3fa3973d3d9d 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanFragment.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanFragment.java @@ -23,7 +23,7 @@ import java.util.List; import java.util.Map; import java.util.Set; -import org.apache.pinot.core.routing.TimeBoundaryInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; import org.apache.pinot.query.planner.PlanFragment; import org.apache.pinot.query.routing.QueryServerInstance; import org.apache.pinot.query.routing.WorkerMetadata; diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanMetadata.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanMetadata.java index ec88a8e4f40d..98690aa86ac7 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanMetadata.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanMetadata.java @@ -27,10 +27,10 @@ import java.util.Map; import java.util.Set; import javax.annotation.Nullable; -import org.apache.pinot.core.routing.TimeBoundaryInfo; +import org.apache.pinot.core.routing.LogicalTableRouteInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; import org.apache.pinot.query.routing.MailboxInfos; import org.apache.pinot.query.routing.QueryServerInstance; -import org.apache.pinot.query.routing.table.LogicalTableRouteInfo; /** diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanVisitor.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanVisitor.java index 848bea8fbc19..77ca8a68e804 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanVisitor.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/DispatchablePlanVisitor.java @@ -24,6 +24,8 @@ import java.util.Set; import org.apache.pinot.calcite.rel.hint.PinotHintOptions; import org.apache.pinot.common.config.provider.TableCache; +import org.apache.pinot.core.routing.LogicalTableRouteInfo; +import org.apache.pinot.core.routing.LogicalTableRouteProvider; import org.apache.pinot.query.planner.plannode.AggregateNode; import org.apache.pinot.query.planner.plannode.ExchangeNode; import org.apache.pinot.query.planner.plannode.ExplainedNode; @@ -39,8 +41,6 @@ import org.apache.pinot.query.planner.plannode.TableScanNode; import org.apache.pinot.query.planner.plannode.ValueNode; import org.apache.pinot.query.planner.plannode.WindowNode; -import org.apache.pinot.query.routing.table.LogicalTableRouteInfo; -import org.apache.pinot.query.routing.table.LogicalTableRouteProvider; public class DispatchablePlanVisitor implements PlanNodeVisitor { diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/TableScanMetadata.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/TableScanMetadata.java index 4eb58c54a3db..2eb05f0f1fee 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/TableScanMetadata.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/TableScanMetadata.java @@ -22,7 +22,7 @@ import java.util.Map; import java.util.Set; import javax.annotation.Nullable; -import org.apache.pinot.core.routing.TimeBoundaryInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; import org.apache.pinot.query.planner.physical.v2.nodes.PhysicalTableScan; diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LeafStageWorkerAssignmentRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LeafStageWorkerAssignmentRule.java index 1ac9d9f1bf75..3461a5388418 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LeafStageWorkerAssignmentRule.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LeafStageWorkerAssignmentRule.java @@ -53,7 +53,7 @@ import org.apache.pinot.core.routing.RoutingTable; import org.apache.pinot.core.routing.ServerRouteInfo; import org.apache.pinot.core.routing.TablePartitionInfo; -import org.apache.pinot.core.routing.TimeBoundaryInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; import org.apache.pinot.core.transport.ServerInstance; import org.apache.pinot.query.context.PhysicalPlannerContext; import org.apache.pinot.query.planner.logical.LeafStageToPinotQuery; diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/StageMetadata.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/StageMetadata.java index 6eca815d6881..74def65a884c 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/StageMetadata.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/StageMetadata.java @@ -20,7 +20,7 @@ import java.util.List; import java.util.Map; -import org.apache.pinot.core.routing.TimeBoundaryInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; import org.apache.pinot.query.planner.physical.DispatchablePlanFragment; diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java index 078b78fc8c49..00e08eef6de9 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java @@ -37,21 +37,21 @@ import org.apache.pinot.calcite.rel.rules.ImmutableTableOptions; import org.apache.pinot.calcite.rel.rules.TableOptions; import org.apache.pinot.common.request.BrokerRequest; +import org.apache.pinot.core.routing.LogicalTableRouteInfo; +import org.apache.pinot.core.routing.LogicalTableRouteProvider; import org.apache.pinot.core.routing.RoutingManager; import org.apache.pinot.core.routing.RoutingTable; import org.apache.pinot.core.routing.ServerRouteInfo; import org.apache.pinot.core.routing.TablePartitionReplicatedServersInfo; -import org.apache.pinot.core.routing.TimeBoundaryInfo; +import org.apache.pinot.core.routing.TableRouteInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; import org.apache.pinot.core.transport.ServerInstance; -import org.apache.pinot.core.transport.TableRouteInfo; import org.apache.pinot.query.planner.PlanFragment; import org.apache.pinot.query.planner.physical.DispatchablePlanContext; import org.apache.pinot.query.planner.physical.DispatchablePlanMetadata; import org.apache.pinot.query.planner.plannode.MailboxSendNode; import org.apache.pinot.query.planner.plannode.PlanNode; import org.apache.pinot.query.planner.plannode.TableScanNode; -import org.apache.pinot.query.routing.table.LogicalTableRouteInfo; -import org.apache.pinot.query.routing.table.LogicalTableRouteProvider; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey; import org.apache.pinot.spi.utils.builder.TableNameBuilder; diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java index f2cd6928ce2d..b4bc0b173d28 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryEnvironmentTestBase.java @@ -32,11 +32,11 @@ import org.apache.commons.collections4.MapUtils; import org.apache.commons.lang3.tuple.Pair; import org.apache.pinot.common.config.provider.TableCache; +import org.apache.pinot.core.routing.MockRoutingManagerFactory; import org.apache.pinot.core.routing.RoutingManager; import org.apache.pinot.core.routing.TablePartitionReplicatedServersInfo; import org.apache.pinot.core.routing.TablePartitionReplicatedServersInfo.PartitionInfo; import org.apache.pinot.query.routing.WorkerManager; -import org.apache.pinot.query.testutils.MockRoutingManagerFactory; import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.data.Schema; import org.apache.pinot.spi.utils.CommonConstants; diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestUtils.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestUtils.java index 572c4c55d79f..73aee721d3ce 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestUtils.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestUtils.java @@ -46,7 +46,7 @@ import org.apache.pinot.core.query.executor.QueryExecutor; import org.apache.pinot.core.query.optimizer.QueryOptimizer; import org.apache.pinot.core.query.request.ServerQueryRequest; -import org.apache.pinot.core.routing.TimeBoundaryInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; import org.apache.pinot.query.planner.plannode.PlanNode; import org.apache.pinot.query.routing.StageMetadata; import org.apache.pinot.query.routing.StagePlan; diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/service/server/QueryServerTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/service/server/QueryServerTest.java index 5cbd5ce6bd8d..96d79b0080b5 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/service/server/QueryServerTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/service/server/QueryServerTest.java @@ -33,7 +33,7 @@ import org.apache.pinot.common.proto.PinotQueryWorkerGrpc; import org.apache.pinot.common.proto.Plan; import org.apache.pinot.common.proto.Worker; -import org.apache.pinot.core.routing.TimeBoundaryInfo; +import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; import org.apache.pinot.query.QueryEnvironment; import org.apache.pinot.query.QueryEnvironmentTestBase; import org.apache.pinot.query.QueryTestSet; diff --git a/pinot-timeseries/pinot-timeseries-planner/src/main/java/org/apache/pinot/tsdb/planner/TimeSeriesQueryEnvironment.java b/pinot-timeseries/pinot-timeseries-planner/src/main/java/org/apache/pinot/tsdb/planner/TimeSeriesQueryEnvironment.java index 69ec67fd72d1..a0246ff57cb8 100644 --- a/pinot-timeseries/pinot-timeseries-planner/src/main/java/org/apache/pinot/tsdb/planner/TimeSeriesQueryEnvironment.java +++ b/pinot-timeseries/pinot-timeseries-planner/src/main/java/org/apache/pinot/tsdb/planner/TimeSeriesQueryEnvironment.java @@ -27,8 +27,8 @@ import java.util.Map; import java.util.Set; import org.apache.pinot.common.config.provider.TableCache; +import org.apache.pinot.core.routing.ImplicitHybridTableRouteProvider; import org.apache.pinot.core.routing.RoutingManager; -import org.apache.pinot.query.routing.table.ImplicitHybridTableRouteProvider; import org.apache.pinot.spi.env.PinotConfiguration; import org.apache.pinot.spi.trace.RequestContext; import org.apache.pinot.tsdb.planner.physical.TableScanVisitor; diff --git a/pinot-timeseries/pinot-timeseries-planner/src/main/java/org/apache/pinot/tsdb/planner/physical/TableScanVisitor.java b/pinot-timeseries/pinot-timeseries-planner/src/main/java/org/apache/pinot/tsdb/planner/physical/TableScanVisitor.java index e00bf4099ae5..cfd7332d88f4 100644 --- a/pinot-timeseries/pinot-timeseries-planner/src/main/java/org/apache/pinot/tsdb/planner/physical/TableScanVisitor.java +++ b/pinot-timeseries/pinot-timeseries-planner/src/main/java/org/apache/pinot/tsdb/planner/physical/TableScanVisitor.java @@ -33,9 +33,9 @@ import org.apache.pinot.common.request.QuerySource; import org.apache.pinot.core.routing.RoutingManager; import org.apache.pinot.core.routing.RoutingTable; +import org.apache.pinot.core.routing.TableRouteInfo; +import org.apache.pinot.core.routing.TableRouteProvider; import org.apache.pinot.core.transport.ServerInstance; -import org.apache.pinot.core.transport.TableRouteInfo; -import org.apache.pinot.query.routing.table.TableRouteProvider; import org.apache.pinot.sql.parsers.CalciteSqlParser; import org.apache.pinot.tsdb.spi.TimeBuckets; import org.apache.pinot.tsdb.spi.plan.BaseTimeSeriesPlanNode; From 9e22441a20f1f6e9cbd9381af1cea512920dae46 Mon Sep 17 00:00:00 2001 From: Himanshu Verma Date: Tue, 29 Jul 2025 15:00:26 +0530 Subject: [PATCH 034/167] Fixed loading the zookeeper browser page is failing (#16440) --- .../resources/app/utils/PinotMethodUtils.ts | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/pinot-controller/src/main/resources/app/utils/PinotMethodUtils.ts b/pinot-controller/src/main/resources/app/utils/PinotMethodUtils.ts index 8f27aab0fae6..e29ffff3583a 100644 --- a/pinot-controller/src/main/resources/app/utils/PinotMethodUtils.ts +++ b/pinot-controller/src/main/resources/app/utils/PinotMethodUtils.ts @@ -737,16 +737,25 @@ const getZookeeperData = (path, count) => { isLeafNode: false, hasChildRendered: true }]; - return getNodeData(path).then((obj)=>{ + + return getNodeData(path).then((obj) => { const { currentNodeData, currentNodeMetadata, currentNodeListStat } = obj; - const pathNames = Object.keys(currentNodeListStat); - pathNames.map((pathName)=>{ + const pathNames = Object.keys(currentNodeListStat || {}); + + pathNames.forEach((pathName) => { + const nodeStat = currentNodeListStat[pathName]; + + // Skip if nodeStat is null or undefined + if (!nodeStat) { + console.warn(`Skipping null node for path: ${pathName}`); + return; + } newTreeData[0].child.push({ nodeId: `${counter++}`, label: pathName, - fullPath: path === '/' ? path+pathName : `${path}/${pathName}`, + fullPath: path === '/' ? path + pathName : `${path}/${pathName}`, child: [], - isLeafNode: currentNodeListStat[pathName].numChildren === 0, + isLeafNode: nodeStat.numChildren === 0, hasChildRendered: false }); }); From 5232e28388d64863ab86df3699257da56af8225c Mon Sep 17 00:00:00 2001 From: Krishan Goyal Date: Tue, 29 Jul 2025 17:00:25 +0530 Subject: [PATCH 035/167] Add logs for segment status checker to be able to debug which segment is failing check (#16437) --- .../helix/SegmentStatusChecker.java | 13 +++++++++--- .../util/ServerQueryInfoFetcher.java | 20 ++++++++++++++++++- .../pinot/core/common/MinionConstants.java | 6 +++++- 3 files changed, 34 insertions(+), 5 deletions(-) diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/SegmentStatusChecker.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/SegmentStatusChecker.java index a86bbaf10a1b..378fb9b55e4c 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/SegmentStatusChecker.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/SegmentStatusChecker.java @@ -366,9 +366,16 @@ private void updateSegmentMetrics(String tableNameWithType, TableConfig tableCon for (Map.Entry entry : stateMap.entrySet()) { String serverInstanceId = entry.getKey(); String segmentState = entry.getValue(); - if ((segmentState.equals(SegmentStateModel.ONLINE) || segmentState.equals(SegmentStateModel.CONSUMING)) - && isServerQueryable(serverQueryInfoFetcher.getServerQueryInfo(serverInstanceId))) { - numEVReplicasUp++; + if (isServerQueryable(serverQueryInfoFetcher.getServerQueryInfo(serverInstanceId))) { + if (segmentState.equals(SegmentStateModel.ONLINE) || segmentState.equals(SegmentStateModel.CONSUMING)) { + numEVReplicasUp++; + } else { + LOGGER.warn("Segment {} in table {} has state {} on instance {}. Marking it as unavailable", + segment, tableNameWithType, segmentState, serverInstanceId); + } + } else { + LOGGER.warn("Segment {} in table {} has state {} on unavailable instance {}. Marking it as unavailable", + segment, tableNameWithType, segmentState, serverInstanceId); } if (segmentState.equals(SegmentStateModel.ERROR)) { errorSegments.add(Pair.of(segment, entry.getKey())); diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/util/ServerQueryInfoFetcher.java b/pinot-controller/src/main/java/org/apache/pinot/controller/util/ServerQueryInfoFetcher.java index 2ac53ae508e3..83a5bc9046b0 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/util/ServerQueryInfoFetcher.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/util/ServerQueryInfoFetcher.java @@ -27,6 +27,8 @@ import org.apache.pinot.controller.helix.core.PinotHelixResourceManager; import org.apache.pinot.spi.utils.CommonConstants; import org.apache.pinot.spi.utils.InstanceTypeUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** @@ -34,6 +36,7 @@ * repeated ZK access. This class is NOT thread-safe. */ public class ServerQueryInfoFetcher { + private static final Logger LOGGER = LoggerFactory.getLogger(ServerQueryInfoFetcher.class); private final PinotHelixResourceManager _pinotHelixResourceManager; private final Map _cache; @@ -51,6 +54,7 @@ public ServerQueryInfo getServerQueryInfo(String instanceId) { private ServerQueryInfo getServerQueryInfoOndemand(String instanceId) { InstanceConfig instanceConfig = _pinotHelixResourceManager.getHelixInstanceConfig(instanceId); if (instanceConfig == null || !InstanceTypeUtils.isServer(instanceId)) { + LOGGER.warn("Instance config for instanceId {} is null or not a server instance", instanceId); return null; } List tags = instanceConfig.getTags(); @@ -59,7 +63,11 @@ private ServerQueryInfo getServerQueryInfoOndemand(String instanceId) { boolean queriesDisabled = record.getBooleanField(CommonConstants.Helix.QUERIES_DISABLED, false); boolean shutdownInProgress = record.getBooleanField(CommonConstants.Helix.IS_SHUTDOWN_IN_PROGRESS, false); - return new ServerQueryInfo(instanceId, tags, null, helixEnabled, queriesDisabled, shutdownInProgress); + ServerQueryInfo queryInfo = new ServerQueryInfo( + instanceId, tags, null, helixEnabled, queriesDisabled, shutdownInProgress); + + LOGGER.info("Fetched ServerQueryInfo for instanceId {}: {}", instanceId, queryInfo); + return queryInfo; } public static class ServerQueryInfo { @@ -91,5 +99,15 @@ public boolean isQueriesDisabled() { public boolean isShutdownInProgress() { return _shutdownInProgress; } + + @Override + public String toString() { + return "ServerQueryInfo{" + + "instanceName='" + _instanceName + '\'' + + ", helixEnabled=" + _helixEnabled + + ", queriesDisabled=" + _queriesDisabled + + ", shutdownInProgress=" + _shutdownInProgress + + '}'; + } } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/common/MinionConstants.java b/pinot-core/src/main/java/org/apache/pinot/core/common/MinionConstants.java index cbb13bd8abd7..4c2285125a3d 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/common/MinionConstants.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/common/MinionConstants.java @@ -31,10 +31,14 @@ private MinionConstants() { public static final String TASK_TIME_SUFFIX = ".time"; public static final String TABLE_NAME_KEY = "tableName"; + + // Input segment name(s) and download url(s) for those segments for the minion task. + // If there are multiple segments, they are separated by SEGMENT_NAME_SEPARATOR. + // The index of the segment name and download url is the same for the same segment. public static final String SEGMENT_NAME_KEY = "segmentName"; public static final String DOWNLOAD_URL_KEY = "downloadURL"; + public static final String UPLOAD_URL_KEY = "uploadURL"; - public static final String DOT_SEPARATOR = "."; public static final String URL_SEPARATOR = ","; public static final String SEGMENT_NAME_SEPARATOR = ","; public static final String AUTH_TOKEN = "authToken"; From b8c8efbf60dda74ec30e771d9ee75dafc4be8048 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 29 Jul 2025 09:47:02 -0700 Subject: [PATCH 036/167] Bump software.amazon.awssdk:bom from 2.32.8 to 2.32.10 (#16446) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 3604138c3353..6c21527cabcb 100644 --- a/pom.xml +++ b/pom.xml @@ -178,7 +178,7 @@ 0.15.1 0.4.7 4.2.2 - 2.32.8 + 2.32.10 1.2.36 1.22.0 2.14.0 From 737029b65410de5c84d92af1d20130cf14506ff9 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 29 Jul 2025 09:47:23 -0700 Subject: [PATCH 037/167] Bump io.grpc:grpc-bom from 1.73.0 to 1.74.0 (#16447) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 6c21527cabcb..b108abc67339 100644 --- a/pom.xml +++ b/pom.xml @@ -239,7 +239,7 @@ 3.25.8 - 1.73.0 + 1.74.0 26.64.0 1.1.1 1.8 From b33a7863e4b0217bf4cbe138fbc978b025129ef1 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 29 Jul 2025 09:47:44 -0700 Subject: [PATCH 038/167] Bump org.webjars:swagger-ui from 5.26.2 to 5.27.0 (#16448) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index b108abc67339..43b7e49313fc 100644 --- a/pom.xml +++ b/pom.xml @@ -151,7 +151,7 @@ 2.47 2.6.1 1.6.16 - 5.26.2 + 5.27.0 3.4.1 2.9.0 2.5.2 From 86808548af8ecf2fea8d343bfbb34cc385291860 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 29 Jul 2025 09:48:03 -0700 Subject: [PATCH 039/167] Bump net.openhft:chronicle-bom from 2.27ea62 to 2.27ea64 (#16449) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 43b7e49313fc..29ef2cd5215f 100644 --- a/pom.xml +++ b/pom.xml @@ -296,7 +296,7 @@ 2.2.0 5.0.4 5.5.1 - 2.27ea62 + 2.27ea64 2.0.6.1 3.9.11 2.2.0 From f9052de934c78d85c7a549197d4079bf017312f6 Mon Sep 17 00:00:00 2001 From: NOOB <43700604+noob-se7en@users.noreply.github.com> Date: Tue, 29 Jul 2025 23:31:46 +0530 Subject: [PATCH 040/167] Enables auto reset of error segments by default (#15736) --- .../main/java/org/apache/pinot/controller/ControllerConf.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/ControllerConf.java b/pinot-controller/src/main/java/org/apache/pinot/controller/ControllerConf.java index 4b64b2abf013..e8ffa8ebc8c3 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/ControllerConf.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/ControllerConf.java @@ -1130,7 +1130,7 @@ public int getSegmentLevelValidationIntervalInSeconds() { public boolean isAutoResetErrorSegmentsOnValidationEnabled() { - return getProperty(ControllerPeriodicTasksConf.AUTO_RESET_ERROR_SEGMENTS_VALIDATION, false); + return getProperty(ControllerPeriodicTasksConf.AUTO_RESET_ERROR_SEGMENTS_VALIDATION, true); } public long getStatusCheckerInitialDelayInSeconds() { From 45389974f6fad5a5f35b074d06063673112e8dc2 Mon Sep 17 00:00:00 2001 From: Shaurya Chaturvedi Date: Tue, 29 Jul 2025 12:09:06 -0700 Subject: [PATCH 041/167] [timeseries] Adding timeseries language endpoint and UI integration (#16424) Co-authored-by: Shaurya Chaturvedi --- .../pinot/controller/ControllerConf.java | 18 ++++++ .../PinotControllerTimeseriesResource.java | 58 +++++++++++++++++++ .../components/Query/TimeseriesQueryPage.tsx | 50 +++++++++++++--- .../src/main/resources/app/requests/index.ts | 3 + 4 files changed, 120 insertions(+), 9 deletions(-) create mode 100644 pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotControllerTimeseriesResource.java diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/ControllerConf.java b/pinot-controller/src/main/java/org/apache/pinot/controller/ControllerConf.java index e8ffa8ebc8c3..67a48e679f05 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/ControllerConf.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/ControllerConf.java @@ -19,6 +19,7 @@ package org.apache.pinot.controller; import com.google.common.base.Preconditions; +import java.util.ArrayList; import java.util.Arrays; import java.util.HashMap; import java.util.List; @@ -27,6 +28,7 @@ import java.util.Random; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; import org.apache.commons.configuration2.Configuration; import org.apache.commons.lang3.StringUtils; import org.apache.helix.controller.rebalancer.strategy.AutoRebalanceStrategy; @@ -37,6 +39,7 @@ import org.apache.pinot.spi.utils.CommonConstants; import org.apache.pinot.spi.utils.Enablement; import org.apache.pinot.spi.utils.TimeUtils; +import org.apache.pinot.tsdb.spi.PinotTimeSeriesConfiguration; import static org.apache.pinot.spi.utils.CommonConstants.Controller.CONFIG_OF_CONTROLLER_METRICS_PREFIX; import static org.apache.pinot.spi.utils.CommonConstants.Controller.CONFIG_OF_INSTANCE_ID; @@ -1333,4 +1336,19 @@ public int getMaxReloadSegmentZkJobs() { public int getMaxForceCommitZkJobs() { return getProperty(CONFIG_OF_MAX_FORCE_COMMIT_JOBS_IN_ZK, ControllerJob.DEFAULT_MAXIMUM_CONTROLLER_JOBS_IN_ZK); } + + /** + * Get the configured timeseries languages from controller configuration. + * @return List of enabled timeseries languages + */ + public List getTimeseriesLanguages() { + String languagesConfig = getProperty(PinotTimeSeriesConfiguration.getEnabledLanguagesConfigKey()); + if (languagesConfig == null) { + return new ArrayList<>(); + } + return Arrays.stream(languagesConfig.split(",")) + .map(String::trim) + .filter(lang -> !lang.isEmpty()) + .collect(Collectors.toList()); + } } diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotControllerTimeseriesResource.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotControllerTimeseriesResource.java new file mode 100644 index 000000000000..74c159c17336 --- /dev/null +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotControllerTimeseriesResource.java @@ -0,0 +1,58 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.controller.api.resources; + +import io.swagger.annotations.ApiOperation; +import java.util.List; +import javax.inject.Inject; +import javax.ws.rs.GET; +import javax.ws.rs.Path; +import javax.ws.rs.Produces; +import javax.ws.rs.core.Context; +import javax.ws.rs.core.HttpHeaders; +import javax.ws.rs.core.MediaType; +import javax.ws.rs.core.Response; +import org.apache.pinot.controller.ControllerConf; +import org.apache.pinot.controller.api.exception.ControllerApplicationException; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +@Path("/timeseries") +public class PinotControllerTimeseriesResource { + public static final Logger LOGGER = LoggerFactory.getLogger(PinotControllerTimeseriesResource.class); + + @Inject + ControllerConf _controllerConf; + + @GET + @Produces(MediaType.APPLICATION_JSON) + @Path("languages") + @ApiOperation(value = "Get timeseries languages from controller configuration", + notes = "Get timeseries languages from controller configuration") + public List getBrokerTimeSeriesLanguages(@Context HttpHeaders headers) { + try { + return _controllerConf.getTimeseriesLanguages(); + } catch (Exception e) { + LOGGER.error("Error fetching timeseries languages from controller configuration", e); + throw new ControllerApplicationException(LOGGER, + "Error fetching timeseries languages from controller configuration: " + e.getMessage(), + Response.Status.INTERNAL_SERVER_ERROR); + } + } +} diff --git a/pinot-controller/src/main/resources/app/components/Query/TimeseriesQueryPage.tsx b/pinot-controller/src/main/resources/app/components/Query/TimeseriesQueryPage.tsx index 964945813beb..7e63d556883b 100644 --- a/pinot-controller/src/main/resources/app/components/Query/TimeseriesQueryPage.tsx +++ b/pinot-controller/src/main/resources/app/components/Query/TimeseriesQueryPage.tsx @@ -38,7 +38,7 @@ import { UnControlled as CodeMirror } from 'react-codemirror2'; import 'codemirror/lib/codemirror.css'; import 'codemirror/theme/material.css'; import 'codemirror/mode/javascript/javascript'; -import { getTimeSeriesQueryResult } from '../../requests'; +import { getTimeSeriesQueryResult, getTimeSeriesLanguages } from '../../requests'; import { useHistory, useLocation } from 'react-router'; import TableToolbar from '../TableToolbar'; import { Resizable } from 're-resizable'; @@ -224,10 +224,6 @@ const jsonoptions = { wordWrap: 'break-word', }; -const SUPPORTED_QUERY_LANGUAGES = [ - { value: 'm3ql', label: 'M3QL' }, -]; - interface TimeseriesQueryConfig { queryLanguage: string; query: string; @@ -252,6 +248,9 @@ const TimeseriesQueryPage = () => { timeout: 60000, }); + const [supportedLanguages, setSupportedLanguages] = useState>([]); + const [languagesLoading, setLanguagesLoading] = useState(true); + const [rawOutput, setRawOutput] = useState(''); const [rawData, setRawData] = useState(null); const [chartSeries, setChartSeries] = useState([]); @@ -264,6 +263,25 @@ const TimeseriesQueryPage = () => { const [selectedMetric, setSelectedMetric] = useState(null); + // Fetch supported languages from controller configuration + useEffect(() => { + const fetchLanguages = async () => { + try { + setLanguagesLoading(true); + const response = await getTimeSeriesLanguages(); + const languages = response.data || []; + + setSupportedLanguages(languages); + } catch (error) { + console.error('Error fetching timeseries languages:', error); + setSupportedLanguages([]); + } finally { + setLanguagesLoading(false); + } + }; + fetchLanguages(); + }, []); + // Update config when URL parameters change useEffect(() => { const urlParams = new URLSearchParams(location.search); @@ -410,6 +428,15 @@ const TimeseriesQueryPage = () => { return ( + {/* Banner for no enabled languages */} + {!languagesLoading && supportedLanguages.length === 0 && ( + + + No timeseries languages enabled. Please configure timeseries languages in your controller, broker and server configurations using the pinot.timeseries.languages property. + + + )} + { lineWrapping: true, indentWithTabs: true, smartIndent: true, + readOnly: supportedLanguages.length === 0, }} className={classes.codeMirror} autoCursor={false} @@ -449,10 +477,11 @@ const TimeseriesQueryPage = () => { @@ -467,6 +496,7 @@ const TimeseriesQueryPage = () => { value={config.startTime} onChange={(e) => handleConfigChange('startTime', e.target.value as string)} placeholder={getOneMinuteAgoTimestamp()} + disabled={supportedLanguages.length === 0} /> @@ -479,6 +509,7 @@ const TimeseriesQueryPage = () => { value={config.endTime} onChange={(e) => handleConfigChange('endTime', e.target.value as string)} placeholder={getCurrentTimestamp()} + disabled={supportedLanguages.length === 0} /> @@ -490,6 +521,7 @@ const TimeseriesQueryPage = () => { type="text" value={config.timeout} onChange={(e) => handleConfigChange('timeout', parseInt(e.target.value as string) || 60000)} + disabled={supportedLanguages.length === 0} /> @@ -499,7 +531,7 @@ const TimeseriesQueryPage = () => { variant="contained" color="primary" onClick={handleExecuteQuery} - disabled={isLoading || !config.query.trim()} + disabled={isLoading || !config.query.trim() || supportedLanguages.length === 0} endIcon={{navigator.platform.includes('Mac') ? '⌘↵' : 'Ctrl+↵'}} > {isLoading ? 'Running Query...' : 'Run Query'} diff --git a/pinot-controller/src/main/resources/app/requests/index.ts b/pinot-controller/src/main/resources/app/requests/index.ts index a069c693dbea..4bef8fdff299 100644 --- a/pinot-controller/src/main/resources/app/requests/index.ts +++ b/pinot-controller/src/main/resources/app/requests/index.ts @@ -238,6 +238,9 @@ export const getQueryResult = (params: Object): Promise export const getTimeSeriesQueryResult = (params: Object): Promise> => transformApi.get(`/timeseries/api/v1/query_range`, { params }); +export const getTimeSeriesLanguages = (): Promise> => + baseApi.get('/timeseries/languages'); + export const getClusterInfo = (): Promise> => baseApi.get('/cluster/info'); From 36ea7c0e26eb80251bb65483391f9f65da486400 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Wed, 30 Jul 2025 18:49:04 +0530 Subject: [PATCH 042/167] Bump org.apache.commons:commons-compress from 1.27.1 to 1.28.0 (#16465) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 29ef2cd5215f..f637be01b8ae 100644 --- a/pom.xml +++ b/pom.xml @@ -205,7 +205,7 @@ 3.18.0 4.5.0 1.14.0 - 1.27.1 + 1.28.0 3.6.1 1.14.0 2.12.0 From 8449763f9d1e44ff6a73d96eae14a33b52a7d657 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Wed, 30 Jul 2025 18:49:12 +0530 Subject: [PATCH 043/167] Bump com.squareup.okio:okio-bom from 3.15.0 to 3.16.0 (#16466) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index f637be01b8ae..7a5e2a8e7253 100644 --- a/pom.xml +++ b/pom.xml @@ -272,7 +272,7 @@ 2.8.3 2.2.0 26.0.2 - 3.15.0 + 3.16.0 2.24.0 3.4 0.10.0 From 16eadb760f7bab7d3850ca7d5dda81e7ed6d1849 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Wed, 30 Jul 2025 18:49:25 +0530 Subject: [PATCH 044/167] Bump net.openhft:chronicle-bom from 2.27ea64 to 2.27ea65 (#16467) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 7a5e2a8e7253..c454ceb8a796 100644 --- a/pom.xml +++ b/pom.xml @@ -296,7 +296,7 @@ 2.2.0 5.0.4 5.5.1 - 2.27ea64 + 2.27ea65 2.0.6.1 3.9.11 2.2.0 From 2314fe6e137f7483a47c4c780d71a18353942841 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Wed, 30 Jul 2025 18:49:37 +0530 Subject: [PATCH 045/167] Bump software.amazon.awssdk:bom from 2.32.10 to 2.32.11 (#16464) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index c454ceb8a796..59302ad4813c 100644 --- a/pom.xml +++ b/pom.xml @@ -178,7 +178,7 @@ 0.15.1 0.4.7 4.2.2 - 2.32.10 + 2.32.11 1.2.36 1.22.0 2.14.0 From 7cfef89d6f7f27c0941fc3877c0aafd618a9eea7 Mon Sep 17 00:00:00 2001 From: Jhow <44998515+J-HowHuang@users.noreply.github.com> Date: Wed, 30 Jul 2025 09:57:40 -0700 Subject: [PATCH 046/167] Incorporate parallelWhitelist and parallelBlacklist to the new Tenant Rebalance API (#16422) --- .../resources/PinotTenantRestletResource.java | 32 ++- .../tenant/DefaultTenantRebalancer.java | 238 +++++++----------- .../rebalance/tenant/TenantRebalancer.java | 34 ++- .../tenant/TenantRebalancerTest.java | 75 +++++- .../tests/TenantRebalanceIntegrationTest.java | 35 --- 5 files changed, 212 insertions(+), 202 deletions(-) diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTenantRestletResource.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTenantRestletResource.java index ba36ff74dddd..fba07869d959 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTenantRestletResource.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTenantRestletResource.java @@ -733,6 +733,18 @@ public TenantRebalanceResult rebalance( + "it will be excluded). Example: table1_REALTIME, table2_REALTIME", example = "") @QueryParam("excludeTables") String excludeTables, + @ApiParam(value = + "Comma separated list of tables with type that are allowed to be rebalanced in parallel. Leaving blank " + + "defaults to " + + "allow all included tables to run in parallel.", + example = "") + @QueryParam("parallelWhitelist") String parallelWhitelist, + @ApiParam(value = + "Comma separated list of tables with type that are restricted to be rebalanced in single thread. These " + + "tables will be removed from parallelWhitelist (that said, if a table appears in both list, " + + "it will be run in single thread).", + example = "") + @QueryParam("parallelBlacklist") String parallelBlacklist, @ApiParam(value = "Show full rebalance results of each table in the response", example = "false") @QueryParam("verboseResult") Boolean verboseResult, @ApiParam(name = "rebalanceConfig", value = "The rebalance config applied to run every table", required = true) @@ -756,17 +768,15 @@ public TenantRebalanceResult rebalance( .map(s -> s.strip().replaceAll("^\"|\"$", "")) .collect(Collectors.toSet())); } - boolean isParallelListSet = !config.getParallelBlacklist().isEmpty() || !config.getParallelWhitelist().isEmpty(); - - boolean isIncludeExcludeListSet = !config.getExcludeTables().isEmpty() || !config.getIncludeTables().isEmpty(); - - // Setting both parallel and include/exclude lists is not allowed. The rebalancer use the old logic if - // parallel white/blacklist is set, otherwise use the new logic (see DefaultTenantRebalancer). Setting both is a - // bad use. - if (isParallelListSet && isIncludeExcludeListSet) { - throw new ControllerApplicationException(LOGGER, - "Bad usage by specifying both include/excludeTables and parallelWhitelist/Blacklist at the same time.", - Response.Status.BAD_REQUEST); + if (parallelBlacklist != null) { + config.setParallelBlacklist(Arrays.stream(StringUtil.split(parallelBlacklist, ',', 0)) + .map(s -> s.strip().replaceAll("^\"|\"$", "")) + .collect(Collectors.toSet())); + } + if (parallelWhitelist != null) { + config.setParallelWhitelist(Arrays.stream(StringUtil.split(parallelWhitelist, ',', 0)) + .map(s -> s.strip().replaceAll("^\"|\"$", "")) + .collect(Collectors.toSet())); } return _tenantRebalancer.rebalance(config); } diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/DefaultTenantRebalancer.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/DefaultTenantRebalancer.java index c7f0d927239b..a23bf2624166 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/DefaultTenantRebalancer.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/DefaultTenantRebalancer.java @@ -19,8 +19,6 @@ package org.apache.pinot.controller.helix.core.rebalance.tenant; import com.google.common.annotations.VisibleForTesting; -import com.google.common.collect.Sets; -import java.util.Deque; import java.util.HashMap; import java.util.HashSet; import java.util.LinkedList; @@ -28,10 +26,11 @@ import java.util.Queue; import java.util.Set; import java.util.UUID; -import java.util.concurrent.ConcurrentLinkedDeque; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.ExecutorService; import java.util.concurrent.atomic.AtomicInteger; +import javax.annotation.Nullable; +import org.apache.commons.lang3.tuple.Pair; import org.apache.pinot.common.exception.TableNotFoundException; import org.apache.pinot.common.utils.config.TagNameUtils; import org.apache.pinot.controller.helix.core.PinotHelixResourceManager; @@ -60,23 +59,20 @@ public DefaultTenantRebalancer(TableRebalanceManager tableRebalanceManager, @Override public TenantRebalanceResult rebalance(TenantRebalanceConfig config) { - if (!config.getParallelWhitelist().isEmpty() || !config.getParallelBlacklist().isEmpty()) { - // If the parallel whitelist or blacklist is set, the old tenant rebalance logic will be used - // TODO: Deprecate the support for this in the future - LOGGER.info("Using the old tenant rebalance logic because parallel whitelist or blacklist is set."); - return rebalanceWithParallelAndSequential(config); - } - return rebalanceWithIncludeExcludeTables(config); - } - - private TenantRebalanceResult rebalanceWithIncludeExcludeTables(TenantRebalanceConfig config) { Map dryRunResults = new HashMap<>(); + + // Step 1: Select the tables to include in this rebalance operation + Set tables = getTenantTables(config.getTenantName()); Set includeTables = config.getIncludeTables(); if (!includeTables.isEmpty()) { tables.retainAll(includeTables); } tables.removeAll(config.getExcludeTables()); + + // Step 2: Dry-run over the selected tables to get the dry-run rebalance results. The result is to be sent as + // response to the user, and their summaries are needed for scheduling the job queue later + tables.forEach(table -> { try { RebalanceConfig rebalanceConfig = RebalanceConfig.copy(config); @@ -89,37 +85,38 @@ private TenantRebalanceResult rebalanceWithIncludeExcludeTables(TenantRebalanceC } }); + // If dry-run was set, return the dry-run results and the job is done here if (config.isDryRun()) { return new TenantRebalanceResult(null, dryRunResults, config.isVerboseResult()); } + // Step 3: Create two queues--parallel and sequential and schedule the tables to these queues based on the + // parallelWhitelist and parallelBlacklist, also their dry-run results. For each table, a job context is created + // and put in the queue for the consuming threads to pick up and run the rebalance operation + String tenantRebalanceJobId = createUniqueRebalanceJobIdentifier(); TenantRebalanceObserver observer = new ZkBasedTenantRebalanceObserver(tenantRebalanceJobId, config.getTenantName(), tables, _pinotHelixResourceManager); observer.onTrigger(TenantRebalanceObserver.Trigger.START_TRIGGER, null, null); - ConcurrentLinkedQueue parallelQueue = createTableQueue(config, dryRunResults); + Pair, Queue> queues = + createParallelAndSequentialQueues(config, dryRunResults, config.getParallelWhitelist(), + config.getParallelBlacklist()); + ConcurrentLinkedQueue parallelQueue = queues.getLeft(); + Queue sequentialQueue = queues.getRight(); + + // Step 4: Spin up threads to consume the parallel queue and sequential queue. + // ensure atleast 1 thread is created to run the sequential table rebalance operations int parallelism = Math.max(config.getDegreeOfParallelism(), 1); + AtomicInteger activeThreads = new AtomicInteger(parallelism); try { for (int i = 0; i < parallelism; i++) { _executorService.submit(() -> { - while (true) { - String table = parallelQueue.poll(); - if (table == null) { - break; - } - RebalanceConfig rebalanceConfig = RebalanceConfig.copy(config); - rebalanceConfig.setDryRun(false); - if (dryRunResults.get(table) - .getRebalanceSummaryResult() - .getSegmentInfo() - .getReplicationFactor() - .getExpectedValueAfterRebalance() == 1) { - rebalanceConfig.setMinAvailableReplicas(0); - } - rebalanceTable(table, rebalanceConfig, dryRunResults.get(table).getJobId(), observer); + doConsumeTablesFromQueue(parallelQueue, config, observer); + if (activeThreads.decrementAndGet() == 0) { + doConsumeTablesFromQueue(sequentialQueue, config, observer); + observer.onSuccess(String.format("Successfully rebalanced tenant %s.", config.getTenantName())); } - observer.onSuccess(String.format("Successfully rebalanced tenant %s.", config.getTenantName())); }); } } catch (Exception exception) { @@ -127,7 +124,9 @@ private TenantRebalanceResult rebalanceWithIncludeExcludeTables(TenantRebalanceC exception.getMessage())); } - // Prepare tenant rebalance result to return + // Step 5: Prepare the rebalance results to be returned to the user. The rebalance jobs are running in the + // background asynchronously. + Map rebalanceResults = new HashMap<>(); for (String table : dryRunResults.keySet()) { RebalanceResult result = dryRunResults.get(table); @@ -143,116 +142,21 @@ private TenantRebalanceResult rebalanceWithIncludeExcludeTables(TenantRebalanceC return new TenantRebalanceResult(tenantRebalanceJobId, rebalanceResults, config.isVerboseResult()); } - // This method implements the old logic for tenant rebalance using parallel whitelist/blacklist. - // Usage of this method would likely be deprecated and not supported in the future. - private TenantRebalanceResult rebalanceWithParallelAndSequential(TenantRebalanceConfig config) { - Map rebalanceResult = new HashMap<>(); - Set tables = getTenantTables(config.getTenantName()); - tables.forEach(table -> { - try { - RebalanceConfig rebalanceConfig = RebalanceConfig.copy(config); - rebalanceConfig.setDryRun(true); - rebalanceResult.put(table, - _tableRebalanceManager.rebalanceTableDryRun(table, rebalanceConfig, createUniqueRebalanceJobIdentifier())); - } catch (TableNotFoundException exception) { - rebalanceResult.put(table, new RebalanceResult(null, RebalanceResult.Status.FAILED, exception.getMessage(), - null, null, null, null, null)); - } - }); - if (config.isDryRun()) { - return new TenantRebalanceResult(null, rebalanceResult, config.isVerboseResult()); - } else { - for (String table : rebalanceResult.keySet()) { - RebalanceResult result = rebalanceResult.get(table); - if (result.getStatus() == RebalanceResult.Status.DONE) { - rebalanceResult.put(table, new RebalanceResult(result.getJobId(), RebalanceResult.Status.IN_PROGRESS, - "In progress, check controller task status for the", result.getInstanceAssignment(), - result.getTierInstanceAssignment(), result.getSegmentAssignment(), result.getPreChecksResult(), - result.getRebalanceSummaryResult())); - } - } - } - - String tenantRebalanceJobId = createUniqueRebalanceJobIdentifier(); - TenantRebalanceObserver observer = new ZkBasedTenantRebalanceObserver(tenantRebalanceJobId, config.getTenantName(), - tables, _pinotHelixResourceManager); - observer.onTrigger(TenantRebalanceObserver.Trigger.START_TRIGGER, null, null); - final Deque sequentialQueue = new LinkedList<>(); - final Deque parallelQueue = new ConcurrentLinkedDeque<>(); - // ensure atleast 1 thread is created to run the sequential table rebalance operations - int parallelism = Math.max(config.getDegreeOfParallelism(), 1); - Set dimTables = getDimensionalTables(config.getTenantName()); - AtomicInteger activeThreads = new AtomicInteger(parallelism); - try { - if (parallelism > 1) { - Set parallelTables; - if (!config.getParallelWhitelist().isEmpty()) { - parallelTables = new HashSet<>(config.getParallelWhitelist()); - } else { - parallelTables = new HashSet<>(tables); - } - if (!config.getParallelBlacklist().isEmpty()) { - parallelTables = Sets.difference(parallelTables, config.getParallelBlacklist()); - } - parallelTables.forEach(table -> { - if (dimTables.contains(table)) { - // prioritise dimension tables - parallelQueue.addFirst(table); - } else { - parallelQueue.addLast(table); - } - }); - Sets.difference(tables, parallelTables).forEach(table -> { - if (dimTables.contains(table)) { - // prioritise dimension tables - sequentialQueue.addFirst(table); - } else { - sequentialQueue.addLast(table); - } - }); - } else { - tables.forEach(table -> { - if (dimTables.contains(table)) { - // prioritise dimension tables - sequentialQueue.addFirst(table); - } else { - sequentialQueue.addLast(table); - } - }); + private void doConsumeTablesFromQueue(Queue queue, RebalanceConfig config, + TenantRebalanceObserver observer) { + while (true) { + TenantTableRebalanceJobContext jobContext = queue.poll(); + if (jobContext == null) { + break; } - - for (int i = 0; i < parallelism; i++) { - _executorService.submit(() -> { - while (true) { - String table = parallelQueue.pollFirst(); - if (table == null) { - break; - } - RebalanceConfig rebalanceConfig = RebalanceConfig.copy(config); - rebalanceConfig.setDryRun(false); - rebalanceTable(table, rebalanceConfig, rebalanceResult.get(table).getJobId(), observer); - } - // Last parallel thread to finish the table rebalance job will pick up the - // sequential table rebalance execution - if (activeThreads.decrementAndGet() == 0) { - RebalanceConfig rebalanceConfig = RebalanceConfig.copy(config); - rebalanceConfig.setDryRun(false); - while (true) { - String table = sequentialQueue.pollFirst(); - if (table == null) { - break; - } - rebalanceTable(table, rebalanceConfig, rebalanceResult.get(table).getJobId(), observer); - } - observer.onSuccess(String.format("Successfully rebalanced tenant %s.", config.getTenantName())); - } - }); + String table = jobContext.getTableName(); + RebalanceConfig rebalanceConfig = RebalanceConfig.copy(config); + rebalanceConfig.setDryRun(false); + if (jobContext.shouldRebalanceWithDowntime()) { + rebalanceConfig.setMinAvailableReplicas(0); } - } catch (Exception exception) { - observer.onError(String.format("Failed to rebalance the tenant %s. Cause: %s", config.getTenantName(), - exception.getMessage())); + rebalanceTable(table, rebalanceConfig, jobContext.getJobId(), observer); } - return new TenantRebalanceResult(tenantRebalanceJobId, rebalanceResult, config.isVerboseResult()); } private Set getDimensionalTables(String tenantName) { @@ -309,14 +213,52 @@ private void rebalanceTable(String tableName, RebalanceConfig config, String reb } } + private static Set getTablesToRunInParallel(Set tables, + @Nullable Set parallelWhitelist, @Nullable Set parallelBlacklist) { + Set parallelTables = new HashSet<>(tables); + if (parallelWhitelist != null && !parallelWhitelist.isEmpty()) { + parallelTables.retainAll(parallelWhitelist); + } + if (parallelBlacklist != null && !parallelBlacklist.isEmpty()) { + parallelTables.removeAll(parallelBlacklist); + } + return parallelTables; + } + + @VisibleForTesting + Pair, Queue> + createParallelAndSequentialQueues( + TenantRebalanceConfig config, Map dryRunResults, @Nullable Set parallelWhitelist, + @Nullable Set parallelBlacklist) { + Set parallelTables = getTablesToRunInParallel(dryRunResults.keySet(), parallelWhitelist, parallelBlacklist); + Map parallelTableDryRunResults = new HashMap<>(); + Map sequentialTableDryRunResults = new HashMap<>(); + dryRunResults.forEach((table, result) -> { + if (parallelTables.contains(table)) { + parallelTableDryRunResults.put(table, result); + } else { + sequentialTableDryRunResults.put(table, result); + } + }); + ConcurrentLinkedQueue parallelQueue = + createTableQueue(config, parallelTableDryRunResults); + Queue sequentialQueue = createTableQueue(config, sequentialTableDryRunResults); + return Pair.of(parallelQueue, sequentialQueue); + } + @VisibleForTesting - ConcurrentLinkedQueue createTableQueue(TenantRebalanceConfig config, + ConcurrentLinkedQueue createTableQueue(TenantRebalanceConfig config, Map dryRunResults) { - Queue firstQueue = new LinkedList<>(); - Queue queue = new LinkedList<>(); - Queue lastQueue = new LinkedList<>(); + Queue firstQueue = new LinkedList<>(); + Queue queue = new LinkedList<>(); + Queue lastQueue = new LinkedList<>(); Set dimTables = getDimensionalTables(config.getTenantName()); - dryRunResults.forEach((table, dryRynResult) -> { + dryRunResults.forEach((table, dryRunResult) -> { + TenantTableRebalanceJobContext jobContext = + new TenantTableRebalanceJobContext(table, dryRunResult.getJobId(), dryRunResult.getRebalanceSummaryResult() + .getSegmentInfo() + .getReplicationFactor() + .getExpectedValueAfterRebalance() == 1); if (dimTables.contains(table)) { // check if the dimension table is a pure scale out or scale in. // pure scale out means that only new servers are added and no servers are removed, vice versa @@ -325,21 +267,21 @@ ConcurrentLinkedQueue createTableQueue(TenantRebalanceConfig config, if (!serverInfo.getServersAdded().isEmpty() && serverInfo.getServersRemoved().isEmpty()) { // dimension table's pure scale OUT should be performed BEFORE other regular tables so that queries involving // joining with dimension table won't fail on the new servers - firstQueue.add(table); + firstQueue.add(jobContext); } else if (serverInfo.getServersAdded().isEmpty() && !serverInfo.getServersRemoved().isEmpty()) { // dimension table's pure scale IN should be performed AFTER other regular tables so that queries involving // joining with dimension table won't fail on the old servers - lastQueue.add(table); + lastQueue.add(jobContext); } else { // the dimension table is not a pure scale out or scale in, which is supposed to be rebalanced manually. // Pre-check should capture and warn about this case. - firstQueue.add(table); + firstQueue.add(jobContext); } } else { - queue.add(table); + queue.add(jobContext); } }); - ConcurrentLinkedQueue tableQueue = new ConcurrentLinkedQueue<>(); + ConcurrentLinkedQueue tableQueue = new ConcurrentLinkedQueue<>(); tableQueue.addAll(firstQueue); tableQueue.addAll(queue); tableQueue.addAll(lastQueue); diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/TenantRebalancer.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/TenantRebalancer.java index 5ce230384d85..951498680677 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/TenantRebalancer.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/TenantRebalancer.java @@ -18,7 +18,39 @@ */ package org.apache.pinot.controller.helix.core.rebalance.tenant; - public interface TenantRebalancer { TenantRebalanceResult rebalance(TenantRebalanceConfig config); + + class TenantTableRebalanceJobContext { + private final String _tableName; + private final String _jobId; + // Whether the rebalance should be done with downtime or minAvailableReplicas=0. + private final boolean _withDowntime; + + /** + * Create a context to run a table rebalance job with in a tenant rebalance operation. + * + * @param tableName The name of the table to rebalance. + * @param jobId The job ID for the rebalance operation. + * @param withDowntime Whether the rebalance should be done with downtime or minAvailableReplicas=0. + * @return The result of the rebalance operation. + */ + public TenantTableRebalanceJobContext(String tableName, String jobId, boolean withDowntime) { + _tableName = tableName; + _jobId = jobId; + _withDowntime = withDowntime; + } + + public String getJobId() { + return _jobId; + } + + public String getTableName() { + return _tableName; + } + + public boolean shouldRebalanceWithDowntime() { + return _withDowntime; + } + } } diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/tenant/TenantRebalancerTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/tenant/TenantRebalancerTest.java index 24c79ebd8b56..0b2f099c53e4 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/tenant/TenantRebalancerTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/tenant/TenantRebalancerTest.java @@ -28,8 +28,10 @@ import java.util.Map; import java.util.Queue; import java.util.Set; +import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import org.apache.commons.lang3.tuple.Pair; import org.apache.pinot.common.assignment.InstancePartitions; import org.apache.pinot.common.tier.TierFactory; import org.apache.pinot.common.utils.config.TagNameUtils; @@ -390,10 +392,15 @@ public void testCreateTableQueue() assertEquals(serverInfo.getServersAdded().size(), 3); assertEquals(serverInfo.getServersRemoved().size(), 0); - Queue tableQueue = tenantRebalancer.createTableQueue(config, dryRunResult.getRebalanceTableResults()); + Queue tableQueue = + tenantRebalancer.createTableQueue(config, dryRunResult.getRebalanceTableResults()); // Dimension Table B should be rebalance first since it is a dim table, and we're doing scale out - assertEquals(tableQueue.poll(), OFFLINE_TABLE_NAME_B); - assertEquals(tableQueue.poll(), OFFLINE_TABLE_NAME_A); + TenantRebalancer.TenantTableRebalanceJobContext jobContext = tableQueue.poll(); + assertNotNull(jobContext); + assertEquals(jobContext.getTableName(), OFFLINE_TABLE_NAME_B); + jobContext = tableQueue.poll(); + assertNotNull(jobContext); + assertEquals(jobContext.getTableName(), OFFLINE_TABLE_NAME_A); // untag server 0, now the rebalance is not a pure scale in/out _helixResourceManager.updateInstanceTags(SERVER_INSTANCE_ID_PREFIX + 0, "", false); @@ -407,8 +414,12 @@ public void testCreateTableQueue() tableQueue = tenantRebalancer.createTableQueue(config, dryRunResult.getRebalanceTableResults()); // Dimension table B should be rebalance first in this case. (it does not matter whether dimension tables are // rebalanced first or last in this case, simply because we defaulted it to be first while non-pure scale in/out) - assertEquals(tableQueue.poll(), OFFLINE_TABLE_NAME_B); - assertEquals(tableQueue.poll(), OFFLINE_TABLE_NAME_A); + jobContext = tableQueue.poll(); + assertNotNull(jobContext); + assertEquals(jobContext.getTableName(), OFFLINE_TABLE_NAME_B); + jobContext = tableQueue.poll(); + assertNotNull(jobContext); + assertEquals(jobContext.getTableName(), OFFLINE_TABLE_NAME_A); // untag the added servers, now the rebalance is a pure scale in for (int i = numServers; i < numServers + numServersToAdd; i++) { @@ -423,8 +434,58 @@ public void testCreateTableQueue() tableQueue = tenantRebalancer.createTableQueue(config, dryRunResult.getRebalanceTableResults()); // Dimension table B should be rebalance last in this case - assertEquals(tableQueue.poll(), OFFLINE_TABLE_NAME_A); - assertEquals(tableQueue.poll(), OFFLINE_TABLE_NAME_B); + jobContext = tableQueue.poll(); + assertNotNull(jobContext); + assertEquals(jobContext.getTableName(), OFFLINE_TABLE_NAME_A); + jobContext = tableQueue.poll(); + assertNotNull(jobContext); + assertEquals(jobContext.getTableName(), OFFLINE_TABLE_NAME_B); + + // set table B in parallel blacklist, so that it ends up in sequential queue, and table A in parallel queue + Pair, + Queue> + queues = + tenantRebalancer.createParallelAndSequentialQueues(config, dryRunResult.getRebalanceTableResults(), null, + Collections.singleton(OFFLINE_TABLE_NAME_B)); + Queue parallelQueue = queues.getLeft(); + Queue sequentialQueue = queues.getRight(); + jobContext = parallelQueue.poll(); + assertNotNull(jobContext); + assertEquals(jobContext.getTableName(), OFFLINE_TABLE_NAME_A); + assertNull(parallelQueue.poll()); + jobContext = sequentialQueue.poll(); + assertNotNull(jobContext); + assertEquals(jobContext.getTableName(), OFFLINE_TABLE_NAME_B); + assertNull(sequentialQueue.poll()); + + // set table B in parallel whitelist, so that it ends up in parallel queue, and table A in sequential queue + queues = tenantRebalancer.createParallelAndSequentialQueues(config, dryRunResult.getRebalanceTableResults(), + Collections.singleton(OFFLINE_TABLE_NAME_B), null); + parallelQueue = queues.getLeft(); + sequentialQueue = queues.getRight(); + jobContext = parallelQueue.poll(); + assertNotNull(jobContext); + assertEquals(jobContext.getTableName(), OFFLINE_TABLE_NAME_B); + assertNull(parallelQueue.poll()); + jobContext = sequentialQueue.poll(); + assertNotNull(jobContext); + assertEquals(jobContext.getTableName(), OFFLINE_TABLE_NAME_A); + assertNull(sequentialQueue.poll()); + + // set both tables in parallel whitelist, and table B in parallel blacklist, so that B ends up in sequential + // queue, and table A in parallel queue + queues = tenantRebalancer.createParallelAndSequentialQueues(config, dryRunResult.getRebalanceTableResults(), + Set.of(OFFLINE_TABLE_NAME_A, OFFLINE_TABLE_NAME_B), Collections.singleton(OFFLINE_TABLE_NAME_B)); + parallelQueue = queues.getLeft(); + sequentialQueue = queues.getRight(); + jobContext = parallelQueue.poll(); + assertNotNull(jobContext); + assertEquals(jobContext.getTableName(), OFFLINE_TABLE_NAME_A); + assertNull(parallelQueue.poll()); + jobContext = sequentialQueue.poll(); + assertNotNull(jobContext); + assertEquals(jobContext.getTableName(), OFFLINE_TABLE_NAME_B); + assertNull(sequentialQueue.poll()); _helixResourceManager.deleteOfflineTable(RAW_TABLE_NAME_A); _helixResourceManager.deleteOfflineTable(RAW_TABLE_NAME_B); diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/TenantRebalanceIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/TenantRebalanceIntegrationTest.java index 7215488e2d8b..3b3f426ee1af 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/TenantRebalanceIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/TenantRebalanceIntegrationTest.java @@ -28,7 +28,6 @@ import static org.testng.Assert.assertFalse; import static org.testng.Assert.assertNotNull; import static org.testng.Assert.assertTrue; -import static org.testng.Assert.fail; public class TenantRebalanceIntegrationTest extends BaseHybridClusterIntegrationTest { @@ -63,40 +62,6 @@ public void testParallelWhitelistBlacklistCompatibility() assertTrue(result.getRebalanceTableResults().containsKey(table2)); } - @Test - public void testParallelWhitelistAndIncludeTablesConflict() - throws Exception { - TenantRebalanceConfig config = new TenantRebalanceConfig(); - config.setTenantName(getServerTenant()); - config.setDryRun(true); - String table1 = getTableName() + "_OFFLINE"; - config.getParallelWhitelist().add(table1); - config.getIncludeTables().add(table1); - - // Test conflict when both are set in the body - try { - sendPostRequest(getRebalanceUrl(), JsonUtils.objectToString(config)); - fail("Expected error when both parallelWhitelist and includeTables are set in body"); - } catch (Exception e) { - assertTrue(e.getMessage() - .contains("Bad usage by specifying both include/excludeTables and parallelWhitelist/Blacklist")); - } - - // Test conflict when parallelWhitelist is set in body and includeTables is set as query param - TenantRebalanceConfig config2 = new TenantRebalanceConfig(); - config2.setTenantName(getServerTenant()); - config2.setDryRun(true); - config2.getParallelWhitelist().add(table1); - String urlWithQuery = getRebalanceUrl() + "?includeTables=" + table1; - try { - sendPostRequest(urlWithQuery, JsonUtils.objectToString(config2)); - fail("Expected error when parallelWhitelist is set in body and includeTables in query param"); - } catch (Exception e) { - assertTrue(e.getMessage() - .contains("Bad usage by specifying both include/excludeTables and parallelWhitelist/Blacklist")); - } - } - @Test public void testIncludeTablesQueryParamOverridesBody() throws Exception { From f0e5eedadcfc3f71dd0fbf2fae172444901b47d2 Mon Sep 17 00:00:00 2001 From: Krishan Goyal Date: Wed, 30 Jul 2025 22:46:39 +0530 Subject: [PATCH 047/167] Reduce log volume added from #16437 (#16459) --- .../helix/SegmentStatusChecker.java | 26 ++++++++++++++----- .../util/ServerQueryInfoFetcher.java | 5 +--- 2 files changed, 21 insertions(+), 10 deletions(-) diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/SegmentStatusChecker.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/SegmentStatusChecker.java index 378fb9b55e4c..9b16f66288ec 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/SegmentStatusChecker.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/SegmentStatusChecker.java @@ -313,6 +313,11 @@ private void updateSegmentMetrics(String tableNameWithType, TableConfig tableCon List partialOnlineSegments = new ArrayList<>(); List segmentsInvalidStartTime = new ArrayList<>(); List segmentsInvalidEndTime = new ArrayList<>(); + + // Track unavailable segments by reason for batched logging + List unavailableSegmentsByState = new ArrayList<>(); + List unavailableSegmentsByInstance = new ArrayList<>(); + for (String segment : segments) { // Number of replicas in ideal state that is in ONLINE/CONSUMING state int numISReplicasUp = 0; @@ -366,16 +371,15 @@ private void updateSegmentMetrics(String tableNameWithType, TableConfig tableCon for (Map.Entry entry : stateMap.entrySet()) { String serverInstanceId = entry.getKey(); String segmentState = entry.getValue(); - if (isServerQueryable(serverQueryInfoFetcher.getServerQueryInfo(serverInstanceId))) { - if (segmentState.equals(SegmentStateModel.ONLINE) || segmentState.equals(SegmentStateModel.CONSUMING)) { + if (segmentState.equals(SegmentStateModel.ONLINE) || segmentState.equals(SegmentStateModel.CONSUMING)) { + if (isServerQueryable(serverQueryInfoFetcher.getServerQueryInfo(serverInstanceId))) { numEVReplicasUp++; } else { - LOGGER.warn("Segment {} in table {} has state {} on instance {}. Marking it as unavailable", - segment, tableNameWithType, segmentState, serverInstanceId); + unavailableSegmentsByInstance.add(segment + + " (state: " + segmentState + " on unavailable " + serverInstanceId + ")"); } } else { - LOGGER.warn("Segment {} in table {} has state {} on unavailable instance {}. Marking it as unavailable", - segment, tableNameWithType, segmentState, serverInstanceId); + unavailableSegmentsByState.add(segment + " (state: " + segmentState + " on " + serverInstanceId + ")"); } if (segmentState.equals(SegmentStateModel.ERROR)) { errorSegments.add(Pair.of(segment, entry.getKey())); @@ -398,6 +402,16 @@ private void updateSegmentMetrics(String tableNameWithType, TableConfig tableCon minEVReplicasUpPercent = Math.min(minEVReplicasUpPercent, numEVReplicasUp * 100 / numISReplicasTotal); } + // Log unavailable segments in batches + if (!unavailableSegmentsByState.isEmpty()) { + LOGGER.warn("Table {} has {} segments marked unavailable due to non-ONLINE/CONSUMING states: {}", + tableNameWithType, unavailableSegmentsByState.size(), logSegments(unavailableSegmentsByState)); + } + if (!unavailableSegmentsByInstance.isEmpty()) { + LOGGER.warn("Table {} has {} segments marked unavailable due to unavailable instances: {}", + tableNameWithType, unavailableSegmentsByInstance.size(), logSegments(unavailableSegmentsByInstance)); + } + if (maxISReplicasUp == Integer.MIN_VALUE) { try { maxISReplicasUp = Math.max(Integer.parseInt(idealState.getReplicas()), 1); diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/util/ServerQueryInfoFetcher.java b/pinot-controller/src/main/java/org/apache/pinot/controller/util/ServerQueryInfoFetcher.java index 83a5bc9046b0..37d13fb4a135 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/util/ServerQueryInfoFetcher.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/util/ServerQueryInfoFetcher.java @@ -63,11 +63,8 @@ private ServerQueryInfo getServerQueryInfoOndemand(String instanceId) { boolean queriesDisabled = record.getBooleanField(CommonConstants.Helix.QUERIES_DISABLED, false); boolean shutdownInProgress = record.getBooleanField(CommonConstants.Helix.IS_SHUTDOWN_IN_PROGRESS, false); - ServerQueryInfo queryInfo = new ServerQueryInfo( + return new ServerQueryInfo( instanceId, tags, null, helixEnabled, queriesDisabled, shutdownInProgress); - - LOGGER.info("Fetched ServerQueryInfo for instanceId {}: {}", instanceId, queryInfo); - return queryInfo; } public static class ServerQueryInfo { From e2f8077ae5f35eb476e4214b2f428be1858ba59a Mon Sep 17 00:00:00 2001 From: "Xiaotian (Jackie) Jiang" <17555551+Jackie-Jiang@users.noreply.github.com> Date: Wed, 30 Jul 2025 13:10:48 -0600 Subject: [PATCH 048/167] Fix var-length dictionary validation on BIG_DECIMAL (#16469) --- .../segment/index/dictionary/DictionaryIndexType.java | 6 +++--- .../pinot/segment/local/utils/TableConfigUtilsTest.java | 8 +++++++- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/dictionary/DictionaryIndexType.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/dictionary/DictionaryIndexType.java index 728f2d65b206..2b359306e1c8 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/dictionary/DictionaryIndexType.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/dictionary/DictionaryIndexType.java @@ -108,9 +108,9 @@ public void validate(FieldIndexConfigs indexConfigs, FieldSpec fieldSpec, TableC DictionaryIndexConfig dictionaryConfig = indexConfigs.getConfig(StandardIndexes.dictionary()); if (dictionaryConfig.isEnabled() && dictionaryConfig.getUseVarLengthDictionary()) { DataType storedType = fieldSpec.getDataType().getStoredType(); - Preconditions.checkState(storedType == DataType.STRING || storedType == DataType.BYTES, - "Cannot create var-length dictionary on column: %s of stored type other than STRING or BYTES", - fieldSpec.getName()); + Preconditions.checkState(!storedType.isFixedWidth(), + "Cannot create var-length dictionary on column: %s of fixed-width stored type: %s", fieldSpec.getName(), + storedType); } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/TableConfigUtilsTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/TableConfigUtilsTest.java index 108fc7783e0e..bc6f1a29cc08 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/TableConfigUtilsTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/TableConfigUtilsTest.java @@ -1441,6 +1441,7 @@ public void testValidateIndexingConfig() { .addSingleValueDimension("myCol", FieldSpec.DataType.STRING) .addSingleValueDimension("bytesCol", FieldSpec.DataType.BYTES) .addSingleValueDimension("intCol", FieldSpec.DataType.INT) + .addSingleValueDimension("bigDecimalCol", FieldSpec.DataType.BIG_DECIMAL) .addMultiValueDimension("multiValCol", FieldSpec.DataType.STRING) .build(); TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(TABLE_NAME) @@ -1707,12 +1708,17 @@ public void testValidateIndexingConfig() { // Expected } + tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(TABLE_NAME) + .setVarLengthDictionaryColumns(Arrays.asList("myCol", "bytesCol", "bigDecimalCol", "multiValCol")) + .build(); + TableConfigUtils.validate(tableConfig, schema); + tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(TABLE_NAME) .setVarLengthDictionaryColumns(Arrays.asList("intCol")) .build(); try { TableConfigUtils.validate(tableConfig, schema); - fail("Should fail for Var length dictionary defined for non string/bytes column"); + fail("Should fail for Var length dictionary defined for fixed-width column"); } catch (Exception e) { // expected } From 510275d39aeadc941c8117e14a007294819fc5ee Mon Sep 17 00:00:00 2001 From: Praveen Date: Wed, 30 Jul 2025 13:20:48 -0700 Subject: [PATCH 049/167] [Query Resource Isolation] Workload Scheduler (#16018) * QRI - WorkloadBudgetManager implementation * Address review comments * scheduler * unit test * review comments: metrics, secondary, resource-manager * remove broker admission * Remove default budget --------- Co-authored-by: Vivek Iyer Vaidyanathan Iyer --- .../pinot/common/metrics/ServerMeter.java | 4 +- .../scheduler/QuerySchedulerFactory.java | 5 + .../query/scheduler/WorkloadScheduler.java | 116 ++++++++++++++++++ .../accounting/WorkloadBudgetManagerTest.java | 26 ++++ .../scheduler/QuerySchedulerFactoryTest.java | 5 + .../accounting/WorkloadBudgetManager.java | 62 ++++++++++ .../pinot/spi/exception/QueryErrorCode.java | 1 + .../pinot/spi/utils/CommonConstants.java | 5 + 8 files changed, 223 insertions(+), 1 deletion(-) create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/WorkloadScheduler.java diff --git a/pinot-common/src/main/java/org/apache/pinot/common/metrics/ServerMeter.java b/pinot-common/src/main/java/org/apache/pinot/common/metrics/ServerMeter.java index 8a73029ac239..f3fd41829b23 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/metrics/ServerMeter.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/metrics/ServerMeter.java @@ -212,7 +212,9 @@ public enum ServerMeter implements AbstractMetrics.Meter { /** * Approximate heap bytes used by the mutable JSON index at the time of index close. */ - MUTABLE_JSON_INDEX_MEMORY_USAGE("bytes", false); + MUTABLE_JSON_INDEX_MEMORY_USAGE("bytes", false), + // Workload Budget exceeded counter + WORKLOAD_BUDGET_EXCEEDED("workloadBudgetExceeded", false, "Number of times workload budget exceeded"); private final String _meterName; private final String _unit; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/QuerySchedulerFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/QuerySchedulerFactory.java index bf6be84be129..ae934d543f06 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/QuerySchedulerFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/QuerySchedulerFactory.java @@ -49,6 +49,7 @@ private QuerySchedulerFactory() { public static final String BINARY_WORKLOAD_ALGORITHM = "binary_workload"; public static final String ALGORITHM_NAME_CONFIG_KEY = "name"; public static final String DEFAULT_QUERY_SCHEDULER_ALGORITHM = FCFS_ALGORITHM; + public static final String WORKLOAD_SCHEDULER_ALGORITHM = "workload"; /** * Static factory to instantiate query scheduler based on scheduler configuration. @@ -83,6 +84,10 @@ public static QueryScheduler create(PinotConfiguration schedulerConfig, QueryExe scheduler = new BinaryWorkloadScheduler(schedulerConfig, queryExecutor, serverMetrics, latestQueryTime, resourceUsageAccountant); break; + case WORKLOAD_SCHEDULER_ALGORITHM: + scheduler = new WorkloadScheduler(schedulerConfig, queryExecutor, serverMetrics, latestQueryTime, + resourceUsageAccountant); + break; default: scheduler = getQuerySchedulerByClassName(schedulerName, schedulerConfig, queryExecutor, serverMetrics, latestQueryTime, diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/WorkloadScheduler.java b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/WorkloadScheduler.java new file mode 100644 index 000000000000..9d0fbb062154 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/WorkloadScheduler.java @@ -0,0 +1,116 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.core.query.scheduler; + +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.ListenableFutureTask; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.atomic.LongAccumulator; +import org.apache.pinot.common.metrics.ServerMeter; +import org.apache.pinot.common.metrics.ServerMetrics; +import org.apache.pinot.common.metrics.ServerQueryPhase; +import org.apache.pinot.common.utils.config.QueryOptionsUtils; +import org.apache.pinot.core.accounting.WorkloadBudgetManager; +import org.apache.pinot.core.query.executor.QueryExecutor; +import org.apache.pinot.core.query.request.ServerQueryRequest; +import org.apache.pinot.core.query.scheduler.resources.QueryExecutorService; +import org.apache.pinot.core.query.scheduler.resources.UnboundedResourceManager; +import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant; +import org.apache.pinot.spi.env.PinotConfiguration; +import org.apache.pinot.spi.query.QueryThreadContext; +import org.apache.pinot.spi.trace.Tracing; +import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.builder.TableNameBuilder; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/** + * Scheduler implementation that supports query admission control based on workload-specific budgets. + * + *

    This class integrates with the {@link WorkloadBudgetManager} to apply CPU and memory budget enforcement + * for different workloads, including primary and secondary workloads.

    + * + *

    Secondary workload configuration is used for queries tagged as "secondary". Queries that exceed their budget + * will be rejected.

    + * + */ +public class WorkloadScheduler extends QueryScheduler { + + private static final Logger LOGGER = LoggerFactory.getLogger(WorkloadScheduler.class); + + private final WorkloadBudgetManager _workloadBudgetManager; + private final ServerMetrics _serverMetrics; + private String _secondaryWorkloadName; + + public WorkloadScheduler(PinotConfiguration config, QueryExecutor queryExecutor, ServerMetrics metrics, + LongAccumulator latestQueryTime, ThreadResourceUsageAccountant resourceUsageAccountant) { + super(config, queryExecutor, new UnboundedResourceManager(config, resourceUsageAccountant), metrics, + latestQueryTime); + _serverMetrics = metrics; + _workloadBudgetManager = Tracing.ThreadAccountantOps.getWorkloadBudgetManager(); + _secondaryWorkloadName = config.getProperty(CommonConstants.Accounting.CONFIG_OF_SECONDARY_WORKLOAD_NAME, + CommonConstants.Accounting.DEFAULT_SECONDARY_WORKLOAD_NAME); + } + + @Override + public String name() { + return "WorkloadScheduler"; + } + + @Override + public ListenableFuture submit(ServerQueryRequest queryRequest) { + if (!_isRunning) { + return shuttingDown(queryRequest); + } + + boolean isSecondary = QueryOptionsUtils.isSecondaryWorkload(queryRequest.getQueryContext().getQueryOptions()); + // This is for backward compatibility, where we want to honor the isSecondary query option to use the secondary + // workload budget. + String workloadName = isSecondary + ? _secondaryWorkloadName + : QueryOptionsUtils.getWorkloadName(queryRequest.getQueryContext().getQueryOptions()); + if (!_workloadBudgetManager.canAdmitQuery(workloadName)) { + // TODO: Explore queuing the query instead of rejecting it. + String tableName = TableNameBuilder.extractRawTableName(queryRequest.getTableNameWithType()); + LOGGER.warn("Workload budget exceeded for workload: {} query: {} table: {}", workloadName, + queryRequest.getRequestId(), tableName); + _serverMetrics.addMeteredValue(workloadName, ServerMeter.WORKLOAD_BUDGET_EXCEEDED, 1L); + _serverMetrics.addMeteredTableValue(tableName, ServerMeter.WORKLOAD_BUDGET_EXCEEDED, 1L); + _serverMetrics.addMeteredGlobalValue(ServerMeter.WORKLOAD_BUDGET_EXCEEDED, 1L); + return outOfCapacity(queryRequest); + } + queryRequest.getTimerContext().startNewPhaseTimer(ServerQueryPhase.SCHEDULER_WAIT); + QueryExecutorService queryExecutorService = _resourceManager.getExecutorService(queryRequest, null); + ExecutorService innerExecutorService = QueryThreadContext.contextAwareExecutorService(queryExecutorService); + ListenableFutureTask queryTask = createQueryFutureTask(queryRequest, innerExecutorService); + _resourceManager.getQueryRunners().submit(queryTask); + return queryTask; + } + + @Override + public void start() { + super.start(); + } + + @Override + public void stop() { + super.stop(); + } +} diff --git a/pinot-core/src/test/java/org/apache/pinot/core/accounting/WorkloadBudgetManagerTest.java b/pinot-core/src/test/java/org/apache/pinot/core/accounting/WorkloadBudgetManagerTest.java index a45bfdf7e8e1..21cc2df4f0bc 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/accounting/WorkloadBudgetManagerTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/accounting/WorkloadBudgetManagerTest.java @@ -115,4 +115,30 @@ void testConcurrentTryChargeSingleWorkload() throws InterruptedException { assertEquals(initialMemBudget - totalMemCharged, remaining._memoryRemaining, "Memory budget mismatch after concurrent updates"); } + + @Test + void testCanAdmitQuery() { + WorkloadBudgetManager manager = new WorkloadBudgetManager(_config); + // Scenario 1: No budget configured -> should admit + assertTrue(manager.canAdmitQuery("unconfigured-workload"), + "Workload without budget should be admitted"); + + // Scenario 2: Budget configured with non-zero remaining -> should admit + String activeWorkload = "active-workload"; + manager.addOrUpdateWorkload(activeWorkload, 100L, 200L); + assertTrue(manager.canAdmitQuery(activeWorkload), "Workload with available budget should be admitted"); + + // Scenario 3: Budget depleted -> should reject + String depletedWorkload = "depleted-workload"; + manager.addOrUpdateWorkload(depletedWorkload, 50L, 50L); + manager.tryCharge(depletedWorkload, 50L, 50L); // deplete + assertFalse(manager.canAdmitQuery(depletedWorkload), + "Workload with depleted budget should be rejected"); + + // Scenario 4: Budget configured with zero cpu remaining -> should reject + String zeroCpuWorkload = "zero-cpu-workload"; + manager.addOrUpdateWorkload(zeroCpuWorkload, 0L, 100L); + assertFalse(manager.canAdmitQuery(zeroCpuWorkload), + "Workload with zero CPU budget should be rejected"); + } } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/QuerySchedulerFactoryTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/QuerySchedulerFactoryTest.java index 9def39b51eeb..47cb2cb3f665 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/QuerySchedulerFactoryTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/scheduler/QuerySchedulerFactoryTest.java @@ -67,6 +67,11 @@ public void testQuerySchedulerFactory() { config.setProperty(QuerySchedulerFactory.ALGORITHM_NAME_CONFIG_KEY, TestQueryScheduler.class.getName()); queryScheduler = QuerySchedulerFactory.create(config, queryExecutor, serverMetrics, latestQueryTime, accountant); assertTrue(queryScheduler instanceof TestQueryScheduler); + + config.setProperty(QuerySchedulerFactory.ALGORITHM_NAME_CONFIG_KEY, + QuerySchedulerFactory.WORKLOAD_SCHEDULER_ALGORITHM); + queryScheduler = QuerySchedulerFactory.create(config, queryExecutor, serverMetrics, latestQueryTime, accountant); + assertTrue(queryScheduler instanceof WorkloadScheduler); } public static final class TestQueryScheduler extends QueryScheduler { diff --git a/pinot-spi/src/main/java/org/apache/pinot/core/accounting/WorkloadBudgetManager.java b/pinot-spi/src/main/java/org/apache/pinot/core/accounting/WorkloadBudgetManager.java index 6e7db555ce6c..5d0004dd628b 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/core/accounting/WorkloadBudgetManager.java +++ b/pinot-spi/src/main/java/org/apache/pinot/core/accounting/WorkloadBudgetManager.java @@ -47,10 +47,41 @@ public WorkloadBudgetManager(PinotConfiguration config) { } _enforcementWindowMs = config.getProperty(CommonConstants.Accounting.CONFIG_OF_WORKLOAD_ENFORCEMENT_WINDOW_MS, CommonConstants.Accounting.DEFAULT_WORKLOAD_ENFORCEMENT_WINDOW_MS); + initSecondaryWorkloadBudget(config); startBudgetResetTask(); LOGGER.info("WorkloadBudgetManager initialized with enforcement window: {}ms", _enforcementWindowMs); } + /** + * This budget is primarily meant to be used for queries that need to be issued in a low priority manner. + * This is fixed budget allocated during host startup and used across all secondary queries. + */ + private void initSecondaryWorkloadBudget(PinotConfiguration config) { + double secondaryCpuPercentage = config.getProperty( + CommonConstants.Accounting.CONFIG_OF_SECONDARY_WORKLOAD_CPU_PERCENTAGE, + CommonConstants.Accounting.DEFAULT_SECONDARY_WORKLOAD_CPU_PERCENTAGE); + + // Don't create a secondary workload if cpu percentage is non-zero. + if (secondaryCpuPercentage <= 0.0) { + return; + } + + String secondaryWorkloadName = config.getProperty( + CommonConstants.Accounting.CONFIG_OF_SECONDARY_WORKLOAD_NAME, + CommonConstants.Accounting.DEFAULT_SECONDARY_WORKLOAD_NAME); + + // The Secondary CPU budget is based on the CPU percentage allocated for secondary workload. + // The memory budget is set to Long.MAX_VALUE for now, since we do not have a specific memory budget for + // secondary queries. + int availableProcessors = Runtime.getRuntime().availableProcessors(); + // Total CPU capacity available in one enforcement window: + // window(ms) × 1_000_000 (ns per ms) × number of logical processors + long totalCpuCapacityNs = _enforcementWindowMs * 1_000_000L * availableProcessors; + long secondaryCpuBudget = (long) (secondaryCpuPercentage * totalCpuCapacityNs); + // TODO: Add memory budget for secondary workload queries + addOrUpdateWorkload(secondaryWorkloadName, secondaryCpuBudget, Long.MAX_VALUE); + } + public void shutdown() { if (!_isEnabled) { return; @@ -145,6 +176,37 @@ private void startBudgetResetTask() { }, _enforcementWindowMs, _enforcementWindowMs, TimeUnit.MILLISECONDS); } + /** + * Determines whether a query for the given workload can be admitted under CPU-only budgets. + * + *

    Admission rules: + *

      + *
    1. If the manager is disabled or no budget exists for the workload, always admit.
      2. + *
    2. If CPU budget remains above zero, admit immediately.
      3. + *
    3. Otherwise, reject (return false).
      4. + *
    + * + *

    Note: This method currently uses a strict check, where CPU and memory budgets must be above zero. + * This may be relaxed in the future to allow for a percentage of other remaining budget to be used. At that point, + * we can have different admission policies like: Strict, Stealing, etc. + * + * @param workload the workload identifier to check budget for + * @return true if the query may be accepted; false if budget is insufficient + */ + public boolean canAdmitQuery(String workload) { + // If disabled or no budget configured, always admit + if (!_isEnabled) { + return true; + } + Budget budget = _workloadBudgets.get(workload); + if (budget == null) { + LOGGER.debug("No budget found for workload: {}", workload); + return true; + } + BudgetStats stats = budget.getStats(); + return stats._cpuRemaining > 0 && stats._memoryRemaining > 0; + } + /** * Represents remaining budget stats. */ diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/exception/QueryErrorCode.java b/pinot-spi/src/main/java/org/apache/pinot/spi/exception/QueryErrorCode.java index c9ca776d1326..a92f5b1be69c 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/exception/QueryErrorCode.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/exception/QueryErrorCode.java @@ -52,6 +52,7 @@ public enum QueryErrorCode { BROKER_REQUEST_SEND(425, "BrokerRequestSend"), SERVER_NOT_RESPONDING(427, "ServerNotResponding"), TOO_MANY_REQUESTS(429, "TooManyRequests"), + WORKLOAD_BUDGET_EXCEEDED(429, "WorkloadBudgetExceededError"), INTERNAL(450, "InternalError"), MERGE_RESPONSE(500, "MergeResponseError"), QUERY_CANCELLATION(503, "QueryCancellationError"), diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java index 4a7676667e56..d8152b1812ca 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java @@ -1600,6 +1600,11 @@ public static class Accounting { public static final int DEFAULT_WORKLOAD_SLEEP_TIME_MS = 1; public static final String DEFAULT_WORKLOAD_NAME = "default"; + public static final String CONFIG_OF_SECONDARY_WORKLOAD_NAME = "accounting.secondary.workload.name"; + public static final String DEFAULT_SECONDARY_WORKLOAD_NAME = "defaultSecondary"; + public static final String CONFIG_OF_SECONDARY_WORKLOAD_CPU_PERCENTAGE = + "accounting.secondary.workload.cpu.percentage"; + public static final double DEFAULT_SECONDARY_WORKLOAD_CPU_PERCENTAGE = 0.0; } public static class ExecutorService { From b6f0319a984a6500d1e2b095d42d7c267b5d8ca1 Mon Sep 17 00:00:00 2001 From: "Xiaotian (Jackie) Jiang" <17555551+Jackie-Jiang@users.noreply.github.com> Date: Wed, 30 Jul 2025 14:57:59 -0600 Subject: [PATCH 050/167] Make VARIANT/UUID non reserved keyword (#16471) --- pinot-common/src/main/codegen/config.fmpp | 21 ++++--- .../sql/parsers/CalciteSqlParserTest.java | 56 ++++++++++++++----- .../function/CastTransformFunction.java | 14 ++--- 3 files changed, 63 insertions(+), 28 deletions(-) diff --git a/pinot-common/src/main/codegen/config.fmpp b/pinot-common/src/main/codegen/config.fmpp index 65de367c2425..29f10c291981 100644 --- a/pinot-common/src/main/codegen/config.fmpp +++ b/pinot-common/src/main/codegen/config.fmpp @@ -44,10 +44,15 @@ data: { keywords: [ "FILE" "ARCHIVE" - "BIG_DECIMAL" - "BYTES" + # Pinot types allowed in CAST function + # LONG - for BIGINT + # BIG_DECIMAL - for DECIMAL + # STRING - for VARCHAR + # BYTES - for VARBINARY "LONG" + "BIG_DECIMAL" "STRING" + "BYTES" ] # List of non-reserved keywords to add; @@ -55,14 +60,16 @@ data: { nonReservedKeywordsToAdd: [ "FILE" "ARCHIVE" - "BIG_DECIMAL" - "BYTES" - "LONG" - "STRING" # Pinot allows using DEFAULT as the catalog name "DEFAULT_" - # Pinot allows using DATETIME as column name + # Pinot allows using LONG/BIG_DECIMAL/STRING/BYTES/DATETIME/VARIANT/UUID as column name + "LONG" + "BIG_DECIMAL" + "STRING" + "BYTES" "DATETIME" + "VARIANT" + "UUID" # The following keywords are reserved in core Calcite, # are reserved in some version of SQL, diff --git a/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlParserTest.java b/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlParserTest.java index e0ec9f18c12d..64479b8c3138 100644 --- a/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlParserTest.java +++ b/pinot-common/src/test/java/org/apache/pinot/sql/parsers/CalciteSqlParserTest.java @@ -18,38 +18,66 @@ */ package org.apache.pinot.sql.parsers; +import org.testng.annotations.DataProvider; import org.testng.annotations.Test; import static org.apache.pinot.sql.parsers.CalciteSqlParser.CALCITE_SQL_PARSER_IDENTIFIER_MAX_LENGTH; import static org.testng.Assert.assertThrows; -import static org.testng.Assert.fail; + public class CalciteSqlParserTest { private static final String SINGLE_CHAR = "a"; + private static final String QUERY_TEMPLATE = "SELECT %s FROM %s"; @Test public void testIdentifierLength() { String tableName = extendIdentifierToMaxLength("exampleTable"); String columnName = extendIdentifierToMaxLength("exampleColumn"); - try { - final String validQuery = "SELECT count(" + columnName + ") FROM " + tableName + " WHERE " + columnName - + " IS NOT NULL"; - CalciteSqlParser.compileToPinotQuery(validQuery); - } catch (Exception ignore) { - // Should not reach this line - fail(); - } - - final String invalidTableNameQuery = "SELECT count(" + columnName + ") FROM " + tableName + SINGLE_CHAR + " WHERE " - + columnName + " IS NOT NULL"; - final String invalidColumnNameQuery = "SELECT count(" + columnName + SINGLE_CHAR + ") FROM " + tableName + " WHERE " - + columnName + SINGLE_CHAR + " IS NOT NULL"; + String validQuery = createQuery(tableName, columnName); + CalciteSqlParser.compileToPinotQuery(validQuery); + + String invalidTableNameQuery = createQuery(columnName, tableName + SINGLE_CHAR); assertThrows(SqlCompilationException.class, () -> CalciteSqlParser.compileToPinotQuery(invalidTableNameQuery)); + String invalidColumnNameQuery = createQuery(columnName + SINGLE_CHAR, tableName); assertThrows(SqlCompilationException.class, () -> CalciteSqlParser.compileToPinotQuery(invalidColumnNameQuery)); } private String extendIdentifierToMaxLength(String identifier) { return identifier + SINGLE_CHAR.repeat(CALCITE_SQL_PARSER_IDENTIFIER_MAX_LENGTH - identifier.length()); } + + private String createQuery(String columnName, String tableName) { + return String.format(QUERY_TEMPLATE, columnName, tableName); + } + + @Test(dataProvider = "nonReservedKeywords") + public void testNonReservedKeywords(String keyword) { + CalciteSqlParser.compileToPinotQuery(createQuery(keyword, "testTable")); + CalciteSqlParser.compileToPinotQuery(createQuery(keyword.toUpperCase(), "testTable")); + } + + @DataProvider + public static Object[][] nonReservedKeywords() { + return new Object[][]{ + new Object[]{"int"}, + new Object[]{"integer"}, + new Object[]{"long"}, + new Object[]{"bigint"}, + new Object[]{"float"}, + new Object[]{"double"}, + new Object[]{"big_decimal"}, + new Object[]{"decimal"}, + new Object[]{"boolean"}, + // TODO: Revisit if we should make "timestamp" non reserved +// new Object[]{"timestamp"}, + new Object[]{"string"}, + new Object[]{"varchar"}, + new Object[]{"bytes"}, + new Object[]{"binary"}, + new Object[]{"varbinary"}, + new Object[]{"variant"}, + new Object[]{"uuid"} + }; + } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/CastTransformFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/CastTransformFunction.java index fd51113f06b2..0c2bd7e1d647 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/CastTransformFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/CastTransformFunction.java @@ -57,10 +57,6 @@ public void init(List arguments, Map c if (castFormatTransformFunction instanceof LiteralTransformFunction) { String targetType = ((LiteralTransformFunction) castFormatTransformFunction).getStringLiteral().toUpperCase(); switch (targetType) { - case "BYTES": - case "VARBINARY": - _resultMetadata = sourceSV ? BYTES_SV_NO_DICTIONARY_METADATA : BYTES_MV_NO_DICTIONARY_METADATA; - break; case "INT": case "INTEGER": _resultMetadata = sourceSV ? INT_SV_NO_DICTIONARY_METADATA : INT_MV_NO_DICTIONARY_METADATA; @@ -93,6 +89,13 @@ public void init(List arguments, Map c case "VARCHAR": _resultMetadata = sourceSV ? STRING_SV_NO_DICTIONARY_METADATA : STRING_MV_NO_DICTIONARY_METADATA; break; + case "JSON": + _resultMetadata = sourceSV ? JSON_SV_NO_DICTIONARY_METADATA : JSON_MV_NO_DICTIONARY_METADATA; + break; + case "BYTES": + case "VARBINARY": + _resultMetadata = sourceSV ? BYTES_SV_NO_DICTIONARY_METADATA : BYTES_MV_NO_DICTIONARY_METADATA; + break; case "INT_ARRAY": case "INTEGER_ARRAY": _resultMetadata = INT_MV_NO_DICTIONARY_METADATA; @@ -110,9 +113,6 @@ public void init(List arguments, Map c case "VARCHAR_ARRAY": _resultMetadata = STRING_MV_NO_DICTIONARY_METADATA; break; - case "JSON": - _resultMetadata = sourceSV ? JSON_SV_NO_DICTIONARY_METADATA : JSON_MV_NO_DICTIONARY_METADATA; - break; default: throw new IllegalArgumentException("Unable to cast expression to type - " + targetType); } From 9c10407e4455a3fa10aa354c73d555c02fa70f9c Mon Sep 17 00:00:00 2001 From: Nihal Jain Date: Thu, 31 Jul 2025 02:28:30 +0530 Subject: [PATCH 051/167] Replace all usages of commons-collections with commons-collections4 (#16468) --- .../pinot/controller/api/resources/PinotTableInstances.java | 2 +- .../helix/core/assignment/segment/SegmentAssignmentUtils.java | 2 +- .../helix/core/realtime/PinotLLCRealtimeSegmentManager.java | 2 +- .../helix/core/rebalance/DefaultRebalancePreChecker.java | 2 +- .../apache/pinot/core/data/manager/BaseTableDataManager.java | 2 +- .../org/apache/pinot/core/routing/LogicalTableRouteInfo.java | 2 +- .../pinot/core/segment/processing/mapper/SegmentMapper.java | 2 +- .../src/main/java/org/apache/pinot/query/QueryEnvironment.java | 2 +- .../pinot/query/runtime/plan/server/ServerPlanRequestUtils.java | 2 +- .../pinot/query/runtime/timeseries/LeafTimeSeriesOperator.java | 2 +- .../local/upsert/merger/PartialUpsertColumnarMerger.java | 2 +- .../pinot/segment/spi/creator/SegmentGeneratorConfig.java | 2 +- .../java/org/apache/pinot/spi/utils/IngestionConfigUtils.java | 2 +- .../pinot/spi/utils/builder/ControllerRequestURLBuilder.java | 2 +- 14 files changed, 14 insertions(+), 14 deletions(-) diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTableInstances.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTableInstances.java index 419b5f1b6f76..c5719e18358b 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTableInstances.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTableInstances.java @@ -50,7 +50,7 @@ import javax.ws.rs.core.HttpHeaders; import javax.ws.rs.core.MediaType; import javax.ws.rs.core.Response; -import org.apache.commons.collections.CollectionUtils; +import org.apache.commons.collections4.CollectionUtils; import org.apache.pinot.common.exception.TableNotFoundException; import org.apache.pinot.common.utils.DatabaseUtils; import org.apache.pinot.common.utils.SimpleHttpResponse; diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/SegmentAssignmentUtils.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/SegmentAssignmentUtils.java index 68ca22fbda49..68ead006bd14 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/SegmentAssignmentUtils.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/SegmentAssignmentUtils.java @@ -30,7 +30,7 @@ import java.util.Set; import java.util.TreeMap; import javax.annotation.Nullable; -import org.apache.commons.collections.CollectionUtils; +import org.apache.commons.collections4.CollectionUtils; import org.apache.helix.HelixManager; import org.apache.helix.store.zk.ZkHelixPropertyStore; import org.apache.helix.zookeeper.datamodel.ZNRecord; diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java index 57cff126e40b..a87a2ea008c3 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java @@ -52,7 +52,7 @@ import java.util.function.BiFunction; import java.util.stream.Collectors; import javax.annotation.Nullable; -import org.apache.commons.collections.CollectionUtils; +import org.apache.commons.collections4.CollectionUtils; import org.apache.commons.lang3.StringUtils; import org.apache.helix.AccessOption; import org.apache.helix.ClusterMessagingService; diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreChecker.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreChecker.java index dc6135a65bb0..c1c360d59e6f 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreChecker.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreChecker.java @@ -28,7 +28,7 @@ import java.util.Set; import java.util.concurrent.ExecutorService; import javax.annotation.Nullable; -import org.apache.commons.collections.CollectionUtils; +import org.apache.commons.collections4.CollectionUtils; import org.apache.hc.client5.http.impl.io.PoolingHttpClientConnectionManager; import org.apache.pinot.common.assignment.InstanceAssignmentConfigUtils; import org.apache.pinot.common.exception.InvalidConfigException; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/BaseTableDataManager.java b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/BaseTableDataManager.java index 79d452e9f092..e4db4f0b50a0 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/BaseTableDataManager.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/BaseTableDataManager.java @@ -44,7 +44,7 @@ import java.util.concurrent.locks.Lock; import javax.annotation.Nullable; import javax.annotation.concurrent.ThreadSafe; -import org.apache.commons.collections.CollectionUtils; +import org.apache.commons.collections4.CollectionUtils; import org.apache.commons.io.FileUtils; import org.apache.commons.lang3.StringUtils; import org.apache.commons.lang3.tuple.Pair; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/routing/LogicalTableRouteInfo.java b/pinot-core/src/main/java/org/apache/pinot/core/routing/LogicalTableRouteInfo.java index 39a3a709c28f..d2dafdd084d7 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/routing/LogicalTableRouteInfo.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/routing/LogicalTableRouteInfo.java @@ -25,7 +25,7 @@ import java.util.Map; import java.util.Set; import javax.annotation.Nullable; -import org.apache.commons.collections.CollectionUtils; +import org.apache.commons.collections4.CollectionUtils; import org.apache.pinot.common.request.BrokerRequest; import org.apache.pinot.common.request.InstanceRequest; import org.apache.pinot.common.request.TableSegmentsInfo; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/segment/processing/mapper/SegmentMapper.java b/pinot-core/src/main/java/org/apache/pinot/core/segment/processing/mapper/SegmentMapper.java index 63a5a0cd14b0..e81b11746868 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/segment/processing/mapper/SegmentMapper.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/segment/processing/mapper/SegmentMapper.java @@ -27,7 +27,7 @@ import java.util.function.Consumer; import java.util.stream.Collectors; import javax.annotation.Nullable; -import org.apache.commons.collections.CollectionUtils; +import org.apache.commons.collections4.CollectionUtils; import org.apache.commons.io.FileUtils; import org.apache.commons.lang3.tuple.Pair; import org.apache.pinot.core.segment.processing.framework.SegmentProcessorConfig; diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java index 1ee92b679cd9..187125bc236e 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/QueryEnvironment.java @@ -54,7 +54,7 @@ import org.apache.calcite.tools.FrameworkConfig; import org.apache.calcite.tools.Frameworks; import org.apache.calcite.tools.RelBuilder; -import org.apache.commons.collections.MapUtils; +import org.apache.commons.collections4.MapUtils; import org.apache.commons.lang3.tuple.Pair; import org.apache.pinot.calcite.rel.rules.PinotImplicitTableHintRule; import org.apache.pinot.calcite.rel.rules.PinotJoinToDynamicBroadcastRule; diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestUtils.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestUtils.java index 73aee721d3ce..c83b7e4358e0 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestUtils.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestUtils.java @@ -29,7 +29,7 @@ import java.util.concurrent.ExecutorService; import java.util.function.BiConsumer; import javax.annotation.Nullable; -import org.apache.commons.collections.MapUtils; +import org.apache.commons.collections4.MapUtils; import org.apache.commons.lang3.tuple.Pair; import org.apache.pinot.common.metrics.ServerMetrics; import org.apache.pinot.common.request.BrokerRequest; diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/timeseries/LeafTimeSeriesOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/timeseries/LeafTimeSeriesOperator.java index bb700cfc394f..427df4f9e412 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/timeseries/LeafTimeSeriesOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/timeseries/LeafTimeSeriesOperator.java @@ -21,7 +21,7 @@ import com.google.common.base.Preconditions; import java.util.Collections; import java.util.concurrent.ExecutorService; -import org.apache.commons.collections.MapUtils; +import org.apache.commons.collections4.MapUtils; import org.apache.pinot.core.operator.blocks.InstanceResponseBlock; import org.apache.pinot.core.operator.blocks.results.AggregationResultsBlock; import org.apache.pinot.core.operator.blocks.results.GroupByResultsBlock; diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/merger/PartialUpsertColumnarMerger.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/merger/PartialUpsertColumnarMerger.java index f310d0f2b573..5a0e97a27201 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/merger/PartialUpsertColumnarMerger.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/upsert/merger/PartialUpsertColumnarMerger.java @@ -21,7 +21,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; -import org.apache.commons.collections.MapUtils; +import org.apache.commons.collections4.MapUtils; import org.apache.pinot.segment.local.segment.readers.LazyRow; import org.apache.pinot.segment.local.upsert.merger.columnar.ForceOverwriteMerger; import org.apache.pinot.segment.local.upsert.merger.columnar.OverwriteMerger; diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/SegmentGeneratorConfig.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/SegmentGeneratorConfig.java index 55f1bc9b499c..eafc92484a8f 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/SegmentGeneratorConfig.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/creator/SegmentGeneratorConfig.java @@ -29,7 +29,7 @@ import java.util.TreeMap; import java.util.concurrent.TimeUnit; import javax.annotation.Nullable; -import org.apache.commons.collections.CollectionUtils; +import org.apache.commons.collections4.CollectionUtils; import org.apache.pinot.segment.spi.creator.name.FixedSegmentNameGenerator; import org.apache.pinot.segment.spi.creator.name.NormalizedDateSegmentNameGenerator; import org.apache.pinot.segment.spi.creator.name.SegmentNameGenerator; diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/IngestionConfigUtils.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/IngestionConfigUtils.java index 6e51b4e7befe..a3841dde9a58 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/IngestionConfigUtils.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/IngestionConfigUtils.java @@ -23,7 +23,7 @@ import java.util.List; import java.util.Map; import java.util.stream.Collectors; -import org.apache.commons.collections.CollectionUtils; +import org.apache.commons.collections4.CollectionUtils; import org.apache.commons.lang3.StringUtils; import org.apache.pinot.spi.config.table.IndexingConfig; import org.apache.pinot.spi.config.table.TableConfig; diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/builder/ControllerRequestURLBuilder.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/builder/ControllerRequestURLBuilder.java index e90e4272398e..d6efa38cb828 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/builder/ControllerRequestURLBuilder.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/builder/ControllerRequestURLBuilder.java @@ -24,7 +24,7 @@ import java.util.Map; import java.util.stream.Collectors; import javax.annotation.Nullable; -import org.apache.commons.collections.CollectionUtils; +import org.apache.commons.collections4.CollectionUtils; import org.apache.commons.lang3.StringUtils; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType; From 46ddb9a7d9dce6c1c4c400c254992b84ff97a63e Mon Sep 17 00:00:00 2001 From: Yash Mayya Date: Thu, 31 Jul 2025 10:56:20 +0530 Subject: [PATCH 052/167] Minor cleanups in logical tables routing related code (#16460) --- .../api/resources/PinotBrokerDebug.java | 10 +-- .../GrpcBrokerRequestHandler.java | 12 ++-- .../broker/routing/BrokerRoutingManager.java | 18 +++--- ...seSingleStageBrokerRequestHandlerTest.java | 4 +- .../reader/PinotServerDataFetcher.scala | 10 +-- .../core/routing/BaseTableRouteInfo.java | 64 ------------------- .../ImplicitHybridTableRouteInfo.java | 46 ++++++------- .../ImplicitHybridTableRouteProvider.java | 12 ++-- .../core/routing/LogicalTableRouteInfo.java | 29 +++------ .../routing/LogicalTableRouteProvider.java | 1 - .../routing/PhysicalTableRouteProvider.java | 9 ++- .../pinot/core/routing/RoutingTable.java | 6 +- ...verRouteInfo.java => SegmentsToQuery.java} | 9 ++- .../pinot/core/routing/TableRouteInfo.java | 57 ++++++++++++++--- .../core/routing/TableRouteProvider.java | 3 +- .../pinot/core/transport/QueryRouter.java | 7 +- .../core/routing/BaseTableRouteTest.java | 8 +-- ...dTableRouteProviderCalculateRouteTest.java | 18 +++--- ...lTableRouteProviderCalculateRouteTest.java | 4 +- .../routing/MockRoutingManagerFactory.java | 6 +- .../core/transport/QueryRoutingTest.java | 12 ++-- .../rules/LeafStageWorkerAssignmentRule.java | 6 +- .../pinot/query/routing/WorkerManager.java | 10 +-- 23 files changed, 156 insertions(+), 205 deletions(-) delete mode 100644 pinot-core/src/main/java/org/apache/pinot/core/routing/BaseTableRouteInfo.java rename pinot-core/src/main/java/org/apache/pinot/core/{transport => routing}/ImplicitHybridTableRouteInfo.java (83%) rename pinot-core/src/main/java/org/apache/pinot/core/routing/{ServerRouteInfo.java => SegmentsToQuery.java} (85%) diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotBrokerDebug.java b/pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotBrokerDebug.java index 19733099b505..331c91f06ee1 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotBrokerDebug.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/api/resources/PinotBrokerDebug.java @@ -55,7 +55,7 @@ import org.apache.pinot.core.auth.ManualAuthorization; import org.apache.pinot.core.auth.TargetType; import org.apache.pinot.core.routing.RoutingTable; -import org.apache.pinot.core.routing.ServerRouteInfo; +import org.apache.pinot.core.routing.SegmentsToQuery; import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; import org.apache.pinot.core.transport.ServerInstance; import org.apache.pinot.core.transport.server.routing.stats.ServerRoutingStatsEntry; @@ -158,11 +158,11 @@ public Map>> getRoutingTable( @ApiResponse(code = 404, message = "Routing not found"), @ApiResponse(code = 500, message = "Internal server error") }) - public Map> getRoutingTableWithOptionalSegments( + public Map> getRoutingTableWithOptionalSegments( @ApiParam(value = "Name of the table") @PathParam("tableName") String tableName, @Context HttpHeaders headers) { tableName = DatabaseUtils.translateTableName(tableName, headers); - Map> result = new TreeMap<>(); + Map> result = new TreeMap<>(); getRoutingTable(tableName, (tableNameWithType, routingTable) -> result.put(tableNameWithType, routingTable.getServerInstanceToSegmentsMap())); if (!result.isEmpty()) { @@ -193,7 +193,7 @@ private void getRoutingTable(String tableName, BiConsumer } private static Map> removeOptionalSegments( - Map serverInstanceToSegmentsMap) { + Map serverInstanceToSegmentsMap) { Map> ret = new HashMap<>(); serverInstanceToSegmentsMap.forEach((k, v) -> ret.put(k, v.getSegments())); return ret; @@ -232,7 +232,7 @@ public Map> getRoutingTableForQuery( @ApiResponse(code = 404, message = "Routing not found"), @ApiResponse(code = 500, message = "Internal server error") }) - public Map getRoutingTableForQueryWithOptionalSegments( + public Map getRoutingTableForQueryWithOptionalSegments( @ApiParam(value = "SQL query (table name should have type suffix)") @QueryParam("query") String query, @Context HttpHeaders httpHeaders) { BrokerRequest brokerRequest = CalciteSqlCompiler.compileToBrokerRequest(query); diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java index 5b9649e92cba..995b92d0ec73 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java @@ -38,7 +38,7 @@ import org.apache.pinot.common.utils.grpc.ServerGrpcQueryClient; import org.apache.pinot.common.utils.grpc.ServerGrpcRequestBuilder; import org.apache.pinot.core.query.reduce.StreamingReduceService; -import org.apache.pinot.core.routing.ServerRouteInfo; +import org.apache.pinot.core.routing.SegmentsToQuery; import org.apache.pinot.core.routing.TableRouteInfo; import org.apache.pinot.core.transport.ServerInstance; import org.apache.pinot.core.transport.ServerRoutingInstance; @@ -92,10 +92,10 @@ protected BrokerResponseNative processBrokerRequest(long requestId, BrokerReques throws Exception { BrokerRequest offlineBrokerRequest = route.getOfflineBrokerRequest(); BrokerRequest realtimeBrokerRequest = route.getRealtimeBrokerRequest(); - // TODO: Routing bases on Map cannot be supported for logical tables. + // TODO: Routing bases on Map cannot be supported for logical tables. // The routing will be replaces to support table to segment list map in the future. - Map offlineRoutingTable = route.getOfflineRoutingTable(); - Map realtimeRoutingTable = route.getRealtimeRoutingTable(); + Map offlineRoutingTable = route.getOfflineRoutingTable(); + Map realtimeRoutingTable = route.getRealtimeRoutingTable(); // TODO: Add servers queried/responded stats assert offlineBrokerRequest != null || realtimeBrokerRequest != null; @@ -121,9 +121,9 @@ protected BrokerResponseNative processBrokerRequest(long requestId, BrokerReques * Query pinot server for data table. */ private void sendRequest(long requestId, TableType tableType, BrokerRequest brokerRequest, - Map routingTable, + Map routingTable, Map> responseMap, boolean trace) { - for (Map.Entry routingEntry : routingTable.entrySet()) { + for (Map.Entry routingEntry : routingTable.entrySet()) { ServerInstance serverInstance = routingEntry.getKey(); // TODO: support optional segments for GrpcQueryServer. List segments = routingEntry.getValue().getSegments(); diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/BrokerRoutingManager.java b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/BrokerRoutingManager.java index d5d345d1abd4..6cb983c0e5a6 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/BrokerRoutingManager.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/BrokerRoutingManager.java @@ -61,7 +61,7 @@ import org.apache.pinot.common.utils.HashUtil; import org.apache.pinot.core.routing.RoutingManager; import org.apache.pinot.core.routing.RoutingTable; -import org.apache.pinot.core.routing.ServerRouteInfo; +import org.apache.pinot.core.routing.SegmentsToQuery; import org.apache.pinot.core.routing.TablePartitionInfo; import org.apache.pinot.core.routing.TablePartitionReplicatedServersInfo; import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; @@ -757,15 +757,15 @@ public RoutingTable getRoutingTable(BrokerRequest brokerRequest, String tableNam selectionResult.getUnavailableSegments(), selectionResult.getNumPrunedSegments()); } - private Map getServerInstanceToSegmentsMap(String tableNameWithType, + private Map getServerInstanceToSegmentsMap(String tableNameWithType, InstanceSelector.SelectionResult selectionResult) { - Map merged = new HashMap<>(); + Map merged = new HashMap<>(); for (Map.Entry entry : selectionResult.getSegmentToInstanceMap().entrySet()) { ServerInstance serverInstance = _enabledServerInstanceMap.get(entry.getValue()); if (serverInstance != null) { - ServerRouteInfo serverRouteInfoInfo = - merged.computeIfAbsent(serverInstance, k -> new ServerRouteInfo(new ArrayList<>(), new ArrayList<>())); - serverRouteInfoInfo.getSegments().add(entry.getKey()); + SegmentsToQuery segmentsToQuery = + merged.computeIfAbsent(serverInstance, k -> new SegmentsToQuery(new ArrayList<>(), new ArrayList<>())); + segmentsToQuery.getSegments().add(entry.getKey()); } else { // Should not happen in normal case unless encountered unexpected exception when updating routing entries _brokerMetrics.addMeteredTableValue(tableNameWithType, BrokerMeter.SERVER_MISSING_FOR_ROUTING, 1L); @@ -774,12 +774,12 @@ private Map getServerInstanceToSegmentsMap(Stri for (Map.Entry entry : selectionResult.getOptionalSegmentToInstanceMap().entrySet()) { ServerInstance serverInstance = _enabledServerInstanceMap.get(entry.getValue()); if (serverInstance != null) { - ServerRouteInfo serverRouteInfo = merged.get(serverInstance); + SegmentsToQuery segmentsToQuery = merged.get(serverInstance); // Skip servers that don't have non-optional segments, so that servers always get some non-optional segments // to process, to be backward compatible. // TODO: allow servers only with optional segments - if (serverRouteInfo != null) { - serverRouteInfo.getOptionalSegments().add(entry.getKey()); + if (segmentsToQuery != null) { + segmentsToQuery.getOptionalSegments().add(entry.getKey()); } } // TODO: Report missing server metrics when we allow servers only with optional segments. diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandlerTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandlerTest.java index 9c315715531f..94b9b7621ece 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandlerTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandlerTest.java @@ -35,7 +35,7 @@ import org.apache.pinot.common.request.PinotQuery; import org.apache.pinot.common.response.broker.BrokerResponseNative; import org.apache.pinot.core.routing.RoutingTable; -import org.apache.pinot.core.routing.ServerRouteInfo; +import org.apache.pinot.core.routing.SegmentsToQuery; import org.apache.pinot.core.routing.TableRouteInfo; import org.apache.pinot.core.transport.ServerInstance; import org.apache.pinot.spi.config.table.TableConfig; @@ -169,7 +169,7 @@ public void testCancelQuery() { when(routingManager.getQueryTimeoutMs(tableName)).thenReturn(10000L); RoutingTable rt = mock(RoutingTable.class); when(rt.getServerInstanceToSegmentsMap()).thenReturn(Map.of(new ServerInstance(new InstanceConfig("server01_9000")), - new ServerRouteInfo(List.of("segment01"), List.of()))); + new SegmentsToQuery(List.of("segment01"), List.of()))); when(routingManager.getRoutingTable(any(), Mockito.anyLong())).thenReturn(rt); QueryQuotaManager queryQuotaManager = mock(QueryQuotaManager.class); when(queryQuotaManager.acquire(anyString())).thenReturn(true); diff --git a/pinot-connectors/pinot-spark-common/src/main/scala/org/apache/pinot/connector/spark/common/reader/PinotServerDataFetcher.scala b/pinot-connectors/pinot-spark-common/src/main/scala/org/apache/pinot/connector/spark/common/reader/PinotServerDataFetcher.scala index 603f25fc0565..f04f47737a0d 100644 --- a/pinot-connectors/pinot-spark-common/src/main/scala/org/apache/pinot/connector/spark/common/reader/PinotServerDataFetcher.scala +++ b/pinot-connectors/pinot-spark-common/src/main/scala/org/apache/pinot/connector/spark/common/reader/PinotServerDataFetcher.scala @@ -24,7 +24,7 @@ import org.apache.pinot.common.metrics.BrokerMetrics import org.apache.pinot.common.request.BrokerRequest import org.apache.pinot.connector.spark.common.partition.PinotSplit import org.apache.pinot.connector.spark.common.{Logging, PinotDataSourceReadOptions, PinotException} -import org.apache.pinot.core.routing.ServerRouteInfo +import org.apache.pinot.core.routing.SegmentsToQuery import org.apache.pinot.core.transport.server.routing.stats.ServerRoutingStatsManager import org.apache.pinot.core.transport.{AsyncQueryResponse, QueryRouter, ServerInstance} import org.apache.pinot.spi.config.table.TableType @@ -93,7 +93,7 @@ private[reader] class PinotServerDataFetcher( dataTables.filter(_.getNumberOfRows > 0) } - private def createRoutingTableForRequest(): JMap[ServerInstance, ServerRouteInfo] = { + private def createRoutingTableForRequest(): JMap[ServerInstance, SegmentsToQuery] = { val nullZkId: String = null val instanceConfig = new InstanceConfig(nullZkId) instanceConfig.setHostName(pinotSplit.serverAndSegments.serverHost) @@ -101,15 +101,15 @@ private[reader] class PinotServerDataFetcher( // TODO: support netty-sec val serverInstance = new ServerInstance(instanceConfig) Map( - serverInstance -> new ServerRouteInfo(pinotSplit.serverAndSegments.segments.asJava, List[String]().asJava) + serverInstance -> new SegmentsToQuery(pinotSplit.serverAndSegments.segments.asJava, List[String]().asJava) ).asJava } private def submitRequestToPinotServer( offlineBrokerRequest: BrokerRequest, - offlineRoutingTable: JMap[ServerInstance, ServerRouteInfo], + offlineRoutingTable: JMap[ServerInstance, SegmentsToQuery], realtimeBrokerRequest: BrokerRequest, - realtimeRoutingTable: JMap[ServerInstance, ServerRouteInfo]): AsyncQueryResponse = { + realtimeRoutingTable: JMap[ServerInstance, SegmentsToQuery]): AsyncQueryResponse = { logInfo(s"Sending request to ${pinotSplit.serverAndSegments.toString}") queryRouter.submitQuery( partitionId, diff --git a/pinot-core/src/main/java/org/apache/pinot/core/routing/BaseTableRouteInfo.java b/pinot-core/src/main/java/org/apache/pinot/core/routing/BaseTableRouteInfo.java deleted file mode 100644 index 43b244d08acf..000000000000 --- a/pinot-core/src/main/java/org/apache/pinot/core/routing/BaseTableRouteInfo.java +++ /dev/null @@ -1,64 +0,0 @@ -/** - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ -package org.apache.pinot.core.routing; - -public abstract class BaseTableRouteInfo implements TableRouteInfo { - - @Override - public boolean isExists() { - return hasOffline() || hasRealtime(); - } - - @Override - public boolean isHybrid() { - return hasOffline() && hasRealtime(); - } - - @Override - public boolean isOffline() { - return hasOffline() && !hasRealtime(); - } - - @Override - public boolean isRealtime() { - return !hasOffline() && hasRealtime(); - } - - @Override - public boolean isRouteExists() { - if (isOffline()) { - return isOfflineRouteExists(); - } else if (isRealtime()) { - return isRealtimeRouteExists(); - } else { - return isOfflineRouteExists() || isRealtimeRouteExists(); - } - } - - @Override - public boolean isDisabled() { - if (isOffline()) { - return isOfflineTableDisabled(); - } else if (isRealtime()) { - return isRealtimeTableDisabled(); - } else { - return isOfflineTableDisabled() && isRealtimeTableDisabled(); - } - } -} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/transport/ImplicitHybridTableRouteInfo.java b/pinot-core/src/main/java/org/apache/pinot/core/routing/ImplicitHybridTableRouteInfo.java similarity index 83% rename from pinot-core/src/main/java/org/apache/pinot/core/transport/ImplicitHybridTableRouteInfo.java rename to pinot-core/src/main/java/org/apache/pinot/core/routing/ImplicitHybridTableRouteInfo.java index 90cc9089a2b8..93a42ebf4852 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/transport/ImplicitHybridTableRouteInfo.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/routing/ImplicitHybridTableRouteInfo.java @@ -16,7 +16,7 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.pinot.core.transport; +package org.apache.pinot.core.routing; import java.util.HashMap; import java.util.List; @@ -26,17 +26,15 @@ import org.apache.commons.collections4.CollectionUtils; import org.apache.pinot.common.request.BrokerRequest; import org.apache.pinot.common.request.InstanceRequest; -import org.apache.pinot.core.routing.BaseTableRouteInfo; -import org.apache.pinot.core.routing.ServerRouteInfo; import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; +import org.apache.pinot.core.transport.ServerInstance; +import org.apache.pinot.core.transport.ServerRoutingInstance; import org.apache.pinot.spi.config.table.QueryConfig; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; -import org.apache.pinot.spi.query.QueryThreadContext; -import org.apache.pinot.spi.utils.CommonConstants; -public class ImplicitHybridTableRouteInfo extends BaseTableRouteInfo { +public class ImplicitHybridTableRouteInfo implements TableRouteInfo { private String _offlineTableName = null; private boolean _isOfflineRouteExists; private TableConfig _offlineTableConfig; @@ -54,16 +52,16 @@ public class ImplicitHybridTableRouteInfo extends BaseTableRouteInfo { private BrokerRequest _offlineBrokerRequest; private BrokerRequest _realtimeBrokerRequest; - private Map _offlineRoutingTable; - private Map _realtimeRoutingTable; + private Map _offlineRoutingTable; + private Map _realtimeRoutingTable; public ImplicitHybridTableRouteInfo() { } public ImplicitHybridTableRouteInfo(@Nullable BrokerRequest offlineBrokerRequest, @Nullable BrokerRequest realtimeBrokerRequest, - @Nullable Map offlineRoutingTable, - @Nullable Map realtimeRoutingTable) { + @Nullable Map offlineRoutingTable, + @Nullable Map realtimeRoutingTable) { _offlineBrokerRequest = offlineBrokerRequest; _realtimeBrokerRequest = realtimeBrokerRequest; _offlineRoutingTable = offlineRoutingTable; @@ -156,21 +154,21 @@ public Set getRealtimeExecutionServers() { @Nullable @Override - public Map getOfflineRoutingTable() { + public Map getOfflineRoutingTable() { return _offlineRoutingTable; } - public void setOfflineRoutingTable(Map offlineRoutingTable) { + public void setOfflineRoutingTable(Map offlineRoutingTable) { _offlineRoutingTable = offlineRoutingTable; } @Nullable @Override - public Map getRealtimeRoutingTable() { + public Map getRealtimeRoutingTable() { return _realtimeRoutingTable; } - public void setRealtimeRoutingTable(Map realtimeRoutingTable) { + public void setRealtimeRoutingTable(Map realtimeRoutingTable) { _realtimeRoutingTable = realtimeRoutingTable; } @@ -272,11 +270,11 @@ public void setNumPrunedSegmentsTotal(int numPrunedSegmentsTotal) { _numPrunedSegmentsTotal = numPrunedSegmentsTotal; } - protected static Map getRequestMapFromRoutingTable(TableType tableType, - Map routingTable, BrokerRequest brokerRequest, long requestId, String brokerId, + private Map getRequestMapFromRoutingTable(TableType tableType, + Map routingTable, BrokerRequest brokerRequest, long requestId, String brokerId, boolean preferTls) { Map requestMap = new HashMap<>(); - for (Map.Entry entry : routingTable.entrySet()) { + for (Map.Entry entry : routingTable.entrySet()) { ServerRoutingInstance serverRoutingInstance = entry.getKey().toServerRoutingInstance(tableType, preferTls); InstanceRequest instanceRequest = getInstanceRequest(requestId, brokerId, brokerRequest, entry.getValue()); requestMap.put(serverRoutingInstance, instanceRequest); @@ -284,18 +282,10 @@ protected static Map getRequestMapFromRo return requestMap; } - protected static InstanceRequest getInstanceRequest(long requestId, String brokerId, BrokerRequest brokerRequest, - ServerRouteInfo segments) { - InstanceRequest instanceRequest = new InstanceRequest(); - instanceRequest.setRequestId(requestId); - instanceRequest.setCid(QueryThreadContext.getCid()); - instanceRequest.setQuery(brokerRequest); - Map queryOptions = brokerRequest.getPinotQuery().getQueryOptions(); - if (queryOptions != null) { - instanceRequest.setEnableTrace(Boolean.parseBoolean(queryOptions.get(CommonConstants.Broker.Request.TRACE))); - } + private InstanceRequest getInstanceRequest(long requestId, String brokerId, BrokerRequest brokerRequest, + SegmentsToQuery segments) { + InstanceRequest instanceRequest = TableRouteInfo.createInstanceRequest(brokerRequest, brokerId, requestId); instanceRequest.setSearchSegments(segments.getSegments()); - instanceRequest.setBrokerId(brokerId); if (CollectionUtils.isNotEmpty(segments.getOptionalSegments())) { // Don't set this field, i.e. leave it as null, if there is no optional segment at all, to be more backward // compatible, as there are places like in multi-stage query engine where this field is not set today when diff --git a/pinot-core/src/main/java/org/apache/pinot/core/routing/ImplicitHybridTableRouteProvider.java b/pinot-core/src/main/java/org/apache/pinot/core/routing/ImplicitHybridTableRouteProvider.java index 14c8e726f81d..1b75067a5e2d 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/routing/ImplicitHybridTableRouteProvider.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/routing/ImplicitHybridTableRouteProvider.java @@ -25,7 +25,6 @@ import org.apache.pinot.common.config.provider.TableCache; import org.apache.pinot.common.request.BrokerRequest; import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; -import org.apache.pinot.core.transport.ImplicitHybridTableRouteInfo; import org.apache.pinot.core.transport.ServerInstance; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.utils.builder.TableNameBuilder; @@ -108,13 +107,12 @@ public void fillRouteMetadata(ImplicitHybridTableRouteInfo tableRouteInfo, Routi @Override public void calculateRoutes(TableRouteInfo tableRouteInfo, RoutingManager routingManager, - BrokerRequest offlineBrokerRequest, - BrokerRequest realtimeBrokerRequest, long requestId) { + BrokerRequest offlineBrokerRequest, BrokerRequest realtimeBrokerRequest, long requestId) { assert (tableRouteInfo.isExists()); String offlineTableName = tableRouteInfo.getOfflineTableName(); String realtimeTableName = tableRouteInfo.getRealtimeTableName(); - Map offlineRoutingTable = null; - Map realtimeRoutingTable = null; + Map offlineRoutingTable = null; + Map realtimeRoutingTable = null; List unavailableSegments = new ArrayList<>(); int numPrunedSegmentsTotal = 0; @@ -128,7 +126,7 @@ public void calculateRoutes(TableRouteInfo tableRouteInfo, RoutingManager routin } if (routingTable != null) { unavailableSegments.addAll(routingTable.getUnavailableSegments()); - Map serverInstanceToSegmentsMap = + Map serverInstanceToSegmentsMap = routingTable.getServerInstanceToSegmentsMap(); if (!serverInstanceToSegmentsMap.isEmpty()) { offlineRoutingTable = serverInstanceToSegmentsMap; @@ -150,7 +148,7 @@ public void calculateRoutes(TableRouteInfo tableRouteInfo, RoutingManager routin } if (routingTable != null) { unavailableSegments.addAll(routingTable.getUnavailableSegments()); - Map serverInstanceToSegmentsMap = + Map serverInstanceToSegmentsMap = routingTable.getServerInstanceToSegmentsMap(); if (!serverInstanceToSegmentsMap.isEmpty()) { realtimeRoutingTable = serverInstanceToSegmentsMap; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/routing/LogicalTableRouteInfo.java b/pinot-core/src/main/java/org/apache/pinot/core/routing/LogicalTableRouteInfo.java index d2dafdd084d7..f4b0170eb9d9 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/routing/LogicalTableRouteInfo.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/routing/LogicalTableRouteInfo.java @@ -31,18 +31,15 @@ import org.apache.pinot.common.request.TableSegmentsInfo; import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; import org.apache.pinot.core.routing.timeboundary.TimeBoundaryStrategy; -import org.apache.pinot.core.transport.ImplicitHybridTableRouteInfo; import org.apache.pinot.core.transport.ServerInstance; import org.apache.pinot.core.transport.ServerRoutingInstance; import org.apache.pinot.spi.config.table.QueryConfig; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; -import org.apache.pinot.spi.query.QueryThreadContext; -import org.apache.pinot.spi.utils.CommonConstants; import org.apache.pinot.spi.utils.builder.TableNameBuilder; -public class LogicalTableRouteInfo extends BaseTableRouteInfo { +public class LogicalTableRouteInfo implements TableRouteInfo { private String _logicalTableName; private List _offlineTables; private List _realtimeTables; @@ -66,7 +63,7 @@ public Map getRequestMap(long requestId, if (_offlineTables != null) { for (TableRouteInfo physicalTableRoute : _offlineTables) { if (physicalTableRoute.getOfflineRoutingTable() != null) { - for (Map.Entry entry : physicalTableRoute.getOfflineRoutingTable() + for (Map.Entry entry : physicalTableRoute.getOfflineRoutingTable() .entrySet()) { TableSegmentsInfo tableSegmentsInfo = new TableSegmentsInfo(); tableSegmentsInfo.setTableName(physicalTableRoute.getOfflineTableName()); @@ -84,7 +81,7 @@ public Map getRequestMap(long requestId, if (_realtimeTables != null) { for (TableRouteInfo physicalTableRoute : _realtimeTables) { if (physicalTableRoute.getRealtimeRoutingTable() != null) { - for (Map.Entry entry : physicalTableRoute.getRealtimeRoutingTable() + for (Map.Entry entry : physicalTableRoute.getRealtimeRoutingTable() .entrySet()) { TableSegmentsInfo tableSegmentsInfo = new TableSegmentsInfo(); tableSegmentsInfo.setTableName(physicalTableRoute.getRealtimeTableName()); @@ -118,16 +115,8 @@ public Map getRequestMap(long requestId, private InstanceRequest getInstanceRequest(long requestId, String brokerId, BrokerRequest brokerRequest, List tableSegmentsInfoList) { - InstanceRequest instanceRequest = new InstanceRequest(); - instanceRequest.setRequestId(requestId); - instanceRequest.setCid(QueryThreadContext.getCid()); - instanceRequest.setQuery(brokerRequest); - Map queryOptions = brokerRequest.getPinotQuery().getQueryOptions(); - if (queryOptions != null) { - instanceRequest.setEnableTrace(Boolean.parseBoolean(queryOptions.get(CommonConstants.Broker.Request.TRACE))); - } + InstanceRequest instanceRequest = TableRouteInfo.createInstanceRequest(brokerRequest, brokerId, requestId); instanceRequest.setTableSegmentsInfoList(tableSegmentsInfoList); - instanceRequest.setBrokerId(brokerId); return instanceRequest; } @@ -182,7 +171,7 @@ public Set getOfflineExecutionServers() { Set offlineExecutionServers = new HashSet<>(); for (TableRouteInfo offlineTable : _offlineTables) { if (offlineTable.isOfflineRouteExists()) { - Map offlineRoutingTable = offlineTable.getOfflineRoutingTable(); + Map offlineRoutingTable = offlineTable.getOfflineRoutingTable(); if (offlineRoutingTable != null) { offlineExecutionServers.addAll(offlineRoutingTable.keySet()); } @@ -199,7 +188,7 @@ public Set getRealtimeExecutionServers() { Set realtimeExecutionServers = new HashSet<>(); for (TableRouteInfo realtimeTable : _realtimeTables) { if (realtimeTable.isRealtimeRouteExists()) { - Map realtimeRoutingTable = realtimeTable.getRealtimeRoutingTable(); + Map realtimeRoutingTable = realtimeTable.getRealtimeRoutingTable(); if (realtimeRoutingTable != null) { realtimeExecutionServers.addAll(realtimeRoutingTable.keySet()); } @@ -212,13 +201,13 @@ public Set getRealtimeExecutionServers() { @Nullable @Override - public Map getOfflineRoutingTable() { + public Map getOfflineRoutingTable() { throw new UnsupportedOperationException(); } @Nullable @Override - public Map getRealtimeRoutingTable() { + public Map getRealtimeRoutingTable() { throw new UnsupportedOperationException(); } @@ -258,6 +247,7 @@ public BrokerRequest getRealtimeBrokerRequest() { return _realtimeBrokerRequest; } + @Override public boolean isOfflineRouteExists() { if (_offlineTables != null) { for (TableRouteInfo offlineTable : _offlineTables) { @@ -269,6 +259,7 @@ public boolean isOfflineRouteExists() { return false; } + @Override public boolean isRealtimeRouteExists() { if (_realtimeTables != null) { for (TableRouteInfo realtimeTable : _realtimeTables) { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/routing/LogicalTableRouteProvider.java b/pinot-core/src/main/java/org/apache/pinot/core/routing/LogicalTableRouteProvider.java index cf33604c382b..db359e404dde 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/routing/LogicalTableRouteProvider.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/routing/LogicalTableRouteProvider.java @@ -26,7 +26,6 @@ import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; import org.apache.pinot.core.routing.timeboundary.TimeBoundaryStrategy; import org.apache.pinot.core.routing.timeboundary.TimeBoundaryStrategyService; -import org.apache.pinot.core.transport.ImplicitHybridTableRouteInfo; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.data.LogicalTableConfig; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/routing/PhysicalTableRouteProvider.java b/pinot-core/src/main/java/org/apache/pinot/core/routing/PhysicalTableRouteProvider.java index b175f97f5a7a..4071dae6453c 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/routing/PhysicalTableRouteProvider.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/routing/PhysicalTableRouteProvider.java @@ -23,7 +23,6 @@ import java.util.List; import java.util.Map; import org.apache.pinot.common.request.BrokerRequest; -import org.apache.pinot.core.transport.ImplicitHybridTableRouteInfo; import org.apache.pinot.core.transport.ServerInstance; @@ -43,8 +42,8 @@ public void calculateRoutes(TableRouteInfo tableRouteInfo, RoutingManager routin BrokerRequest offlineBrokerRequest, BrokerRequest realtimeBrokerRequest, long requestId) { String offlineTableName = tableRouteInfo.getOfflineTableName(); String realtimeTableName = tableRouteInfo.getRealtimeTableName(); - Map offlineRoutingTable = null; - Map realtimeRoutingTable = null; + Map offlineRoutingTable = null; + Map realtimeRoutingTable = null; List unavailableSegments = new ArrayList<>(); int numPrunedSegmentsTotal = 0; @@ -59,7 +58,7 @@ public void calculateRoutes(TableRouteInfo tableRouteInfo, RoutingManager routin } if (routingTable != null) { unavailableSegments.addAll(routingTable.getUnavailableSegments()); - Map serverInstanceToSegmentsMap = + Map serverInstanceToSegmentsMap = routingTable.getServerInstanceToSegmentsMap(); if (!serverInstanceToSegmentsMap.isEmpty()) { offlineRoutingTable = serverInstanceToSegmentsMap; @@ -82,7 +81,7 @@ public void calculateRoutes(TableRouteInfo tableRouteInfo, RoutingManager routin } if (routingTable != null) { unavailableSegments.addAll(routingTable.getUnavailableSegments()); - Map serverInstanceToSegmentsMap = + Map serverInstanceToSegmentsMap = routingTable.getServerInstanceToSegmentsMap(); if (!serverInstanceToSegmentsMap.isEmpty()) { realtimeRoutingTable = serverInstanceToSegmentsMap; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/routing/RoutingTable.java b/pinot-core/src/main/java/org/apache/pinot/core/routing/RoutingTable.java index 5a7407805d6d..a4c3bd3ddbf2 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/routing/RoutingTable.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/routing/RoutingTable.java @@ -28,18 +28,18 @@ public class RoutingTable { // the newly created consuming segments. Such segments were simply skipped by brokers at query routing time, but that // had caused wrong query results, particularly for upsert tables. Instead, we should pass such segments to servers // and let them decide how to handle them, e.g. skip them upon issues or include them for better query results. - private final Map _serverInstanceToSegmentsMap; + private final Map _serverInstanceToSegmentsMap; private final List _unavailableSegments; private final int _numPrunedSegments; - public RoutingTable(Map serverInstanceToSegmentsMap, + public RoutingTable(Map serverInstanceToSegmentsMap, List unavailableSegments, int numPrunedSegments) { _serverInstanceToSegmentsMap = serverInstanceToSegmentsMap; _unavailableSegments = unavailableSegments; _numPrunedSegments = numPrunedSegments; } - public Map getServerInstanceToSegmentsMap() { + public Map getServerInstanceToSegmentsMap() { return _serverInstanceToSegmentsMap; } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/routing/ServerRouteInfo.java b/pinot-core/src/main/java/org/apache/pinot/core/routing/SegmentsToQuery.java similarity index 85% rename from pinot-core/src/main/java/org/apache/pinot/core/routing/ServerRouteInfo.java rename to pinot-core/src/main/java/org/apache/pinot/core/routing/SegmentsToQuery.java index cd2b52053bee..3acbf37594d0 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/routing/ServerRouteInfo.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/routing/SegmentsToQuery.java @@ -21,20 +21,19 @@ import java.util.List; /** - * Class representing the route information for a server. - * It contains the list of segments and optional segments assigned to the server. + * Wrapper class around the list of segments and optional segments that need to be queried on a particular server. */ -public class ServerRouteInfo { +public class SegmentsToQuery { private final List _segments; private final List _optionalSegments; /** - * Constructor for ServerRouteInfo. + * Constructor for SegmentsToQuery. * * @param segments List of segments assigned to the server. * @param optionalSegments List of optional segments assigned to the server. */ - public ServerRouteInfo(List segments, List optionalSegments) { + public SegmentsToQuery(List segments, List optionalSegments) { _segments = segments; _optionalSegments = optionalSegments; } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/routing/TableRouteInfo.java b/pinot-core/src/main/java/org/apache/pinot/core/routing/TableRouteInfo.java index cf8220cbb2a2..9be811dd6ac7 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/routing/TableRouteInfo.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/routing/TableRouteInfo.java @@ -29,6 +29,8 @@ import org.apache.pinot.core.transport.ServerRoutingInstance; import org.apache.pinot.spi.config.table.QueryConfig; import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.query.QueryThreadContext; +import org.apache.pinot.spi.utils.CommonConstants; /** @@ -110,7 +112,7 @@ public interface TableRouteInfo { * @return a map of server instances to their route information for the offline table, or null if not available */ @Nullable - Map getOfflineRoutingTable(); + Map getOfflineRoutingTable(); /** * Gets the routing table for the realtime table, if available. @@ -120,35 +122,43 @@ public interface TableRouteInfo { * @return a map of server instances to their route information for the realtime table, or null if not available */ @Nullable - Map getRealtimeRoutingTable(); + Map getRealtimeRoutingTable(); /** * Checks if the table exists. A table exists if all required TableConfig objects are available. * * @return true if the table exists, false otherwise */ - boolean isExists(); + default boolean isExists() { + return hasOffline() || hasRealtime(); + } /** * Checks if there are both offline and realtime tables. * * @return true if the table is a hybrid table, false otherwise */ - boolean isHybrid(); + default boolean isHybrid() { + return hasOffline() && hasRealtime(); + } /** * Checks if the table has offline tables only * * @return true if the table has offline tables only, false otherwise */ - boolean isOffline(); + default boolean isOffline() { + return hasOffline() && !hasRealtime(); + } /** * Checks if the table has realtime tables only. * - * @return true if the table table has realtime tables only, false otherwise + * @return true if the table has realtime tables only, false otherwise */ - boolean isRealtime(); + default boolean isRealtime() { + return !hasOffline() && hasRealtime(); + } /** * Checks if the table has at least 1 offline table. It may or may not have realtime tables as well. @@ -169,7 +179,15 @@ public interface TableRouteInfo { * * @return true if any route exists, false otherwise */ - boolean isRouteExists(); + default boolean isRouteExists() { + if (isOffline()) { + return isOfflineRouteExists(); + } else if (isRealtime()) { + return isRealtimeRouteExists(); + } else { + return isOfflineRouteExists() || isRealtimeRouteExists(); + } + } boolean isOfflineRouteExists(); @@ -180,7 +198,15 @@ public interface TableRouteInfo { * * @return true if the table is disabled, false otherwise */ - boolean isDisabled(); + default boolean isDisabled() { + if (isOffline()) { + return isOfflineTableDisabled(); + } else if (isRealtime()) { + return isRealtimeTableDisabled(); + } else { + return isOfflineTableDisabled() && isRealtimeTableDisabled(); + } + } boolean isOfflineTableDisabled(); @@ -215,4 +241,17 @@ public interface TableRouteInfo { * @return A map of ServerRoutingInstance and InstanceRequest */ Map getRequestMap(long requestId, String brokerId, boolean preferTls); + + static InstanceRequest createInstanceRequest(BrokerRequest brokerRequest, String brokerId, long requestId) { + InstanceRequest instanceRequest = new InstanceRequest(); + instanceRequest.setBrokerId(brokerId); + instanceRequest.setRequestId(requestId); + instanceRequest.setCid(QueryThreadContext.getCid()); + instanceRequest.setQuery(brokerRequest); + Map queryOptions = brokerRequest.getPinotQuery().getQueryOptions(); + if (queryOptions != null) { + instanceRequest.setEnableTrace(Boolean.parseBoolean(queryOptions.get(CommonConstants.Broker.Request.TRACE))); + } + return instanceRequest; + } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/routing/TableRouteProvider.java b/pinot-core/src/main/java/org/apache/pinot/core/routing/TableRouteProvider.java index 22912b7f4ddd..f4c7cb39a1b2 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/routing/TableRouteProvider.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/routing/TableRouteProvider.java @@ -27,8 +27,7 @@ * of the metadata are table config, broker routing information and the broker request. */ public interface TableRouteProvider { - TableRouteInfo getTableRouteInfo(String tableName, TableCache tableCache, - RoutingManager routingManager); + TableRouteInfo getTableRouteInfo(String tableName, TableCache tableCache, RoutingManager routingManager); /** * Calculate the Routing Table for a query. The routing table consists of the server name and list of segments that diff --git a/pinot-core/src/main/java/org/apache/pinot/core/transport/QueryRouter.java b/pinot-core/src/main/java/org/apache/pinot/core/transport/QueryRouter.java index a01f248751ef..fd77b1b5e018 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/transport/QueryRouter.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/transport/QueryRouter.java @@ -32,7 +32,8 @@ import org.apache.pinot.common.request.BrokerRequest; import org.apache.pinot.common.request.InstanceRequest; import org.apache.pinot.common.utils.config.QueryOptionsUtils; -import org.apache.pinot.core.routing.ServerRouteInfo; +import org.apache.pinot.core.routing.ImplicitHybridTableRouteInfo; +import org.apache.pinot.core.routing.SegmentsToQuery; import org.apache.pinot.core.routing.TableRouteInfo; import org.apache.pinot.core.transport.server.routing.stats.ServerRoutingStatsManager; import org.apache.pinot.spi.config.table.TableType; @@ -86,9 +87,9 @@ public QueryRouter(String brokerId, BrokerMetrics brokerMetrics, @Nullable Netty public AsyncQueryResponse submitQuery(long requestId, String rawTableName, @Nullable BrokerRequest offlineBrokerRequest, - @Nullable Map offlineRoutingTable, + @Nullable Map offlineRoutingTable, @Nullable BrokerRequest realtimeBrokerRequest, - @Nullable Map realtimeRoutingTable, long timeoutMs) { + @Nullable Map realtimeRoutingTable, long timeoutMs) { TableRouteInfo tableRouteInfo = new ImplicitHybridTableRouteInfo(offlineBrokerRequest, realtimeBrokerRequest, offlineRoutingTable, realtimeRoutingTable); diff --git a/pinot-core/src/test/java/org/apache/pinot/core/routing/BaseTableRouteTest.java b/pinot-core/src/test/java/org/apache/pinot/core/routing/BaseTableRouteTest.java index c3d70431a74d..d9a274d097d0 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/routing/BaseTableRouteTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/routing/BaseTableRouteTest.java @@ -439,12 +439,12 @@ protected TableRouteInfo getLogicalTableRouteInfo(String tableName, String logic return _logicalTableRouteProvider.getTableRouteInfo(logicalTableName, _tableCache, _routingManager); } - static void assertRoutingTableEqual(Map routeComputer, + static void assertRoutingTableEqual(Map routeComputer, Map> expectedRealtimeRoutingTable) { - for (Map.Entry entry : routeComputer.entrySet()) { + for (Map.Entry entry : routeComputer.entrySet()) { ServerInstance serverInstance = entry.getKey(); - ServerRouteInfo serverRouteInfo = entry.getValue(); - Set segments = ImmutableSet.copyOf(serverRouteInfo.getSegments()); + SegmentsToQuery segmentsToQuery = entry.getValue(); + Set segments = ImmutableSet.copyOf(segmentsToQuery.getSegments()); assertTrue(expectedRealtimeRoutingTable.containsKey(serverInstance.toString())); assertEquals(expectedRealtimeRoutingTable.get(serverInstance.toString()), segments); } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/routing/ImplicitHybridTableRouteProviderCalculateRouteTest.java b/pinot-core/src/test/java/org/apache/pinot/core/routing/ImplicitHybridTableRouteProviderCalculateRouteTest.java index e9cb15fbba09..a7bddc5ee4e6 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/routing/ImplicitHybridTableRouteProviderCalculateRouteTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/routing/ImplicitHybridTableRouteProviderCalculateRouteTest.java @@ -142,16 +142,16 @@ void testDisabledTable(String tableName) { private static class GetTableRouteResult { public final BrokerRequest _offlineBrokerRequest; public final BrokerRequest _realtimeBrokerRequest; - public final Map _offlineRoutingTable; - public final Map _realtimeRoutingTable; + public final Map _offlineRoutingTable; + public final Map _realtimeRoutingTable; public final List _unavailableSegments; public final int _numPrunedSegmentsTotal; public final boolean _offlineTableDisabled; public final boolean _realtimeTableDisabled; public GetTableRouteResult(BrokerRequest offlineBrokerRequest, BrokerRequest realtimeBrokerRequest, - Map offlineRoutingTable, - Map realtimeRoutingTable, List unavailableSegments, + Map offlineRoutingTable, + Map realtimeRoutingTable, List unavailableSegments, int numPrunedSegmentsTotal, boolean offlineTableDisabled, boolean realtimeTableDisabled) { _offlineBrokerRequest = offlineBrokerRequest; _realtimeBrokerRequest = realtimeBrokerRequest; @@ -202,8 +202,8 @@ private static GetTableRouteResult getTableRouting(String tableName, RoutingMana } } - Map offlineRoutingTable = null; - Map realtimeRoutingTable = null; + Map offlineRoutingTable = null; + Map realtimeRoutingTable = null; BrokerRequestPair brokerRequestPair = getBrokerRequestPair(tableName, offlineTableName != null, realtimeTableName != null, offlineTableName, realtimeTableName); @@ -224,7 +224,7 @@ private static GetTableRouteResult getTableRouting(String tableName, RoutingMana } if (routingTable != null) { unavailableSegments.addAll(routingTable.getUnavailableSegments()); - Map serverInstanceToSegmentsMap = + Map serverInstanceToSegmentsMap = routingTable.getServerInstanceToSegmentsMap(); if (!serverInstanceToSegmentsMap.isEmpty()) { offlineRoutingTable = serverInstanceToSegmentsMap; @@ -245,7 +245,7 @@ private static GetTableRouteResult getTableRouting(String tableName, RoutingMana } if (routingTable != null) { unavailableSegments.addAll(routingTable.getUnavailableSegments()); - Map serverInstanceToSegmentsMap = + Map serverInstanceToSegmentsMap = routingTable.getServerInstanceToSegmentsMap(); if (!serverInstanceToSegmentsMap.isEmpty()) { realtimeRoutingTable = serverInstanceToSegmentsMap; @@ -266,7 +266,7 @@ private static GetTableRouteResult getTableRouting(String tableName, RoutingMana /** * Checks if two table routes are the same. A expected routingTable is a Map> where the key is the * server name and the value is a set of segments. This is compared to the routing table - * Map + * Map * @param tableName * @param expectedOfflineRoutingTable * @param expectedRealtimeRoutingTable diff --git a/pinot-core/src/test/java/org/apache/pinot/core/routing/LogicalTableRouteProviderCalculateRouteTest.java b/pinot-core/src/test/java/org/apache/pinot/core/routing/LogicalTableRouteProviderCalculateRouteTest.java index 70311857fb1b..1b943c75b19c 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/routing/LogicalTableRouteProviderCalculateRouteTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/routing/LogicalTableRouteProviderCalculateRouteTest.java @@ -66,7 +66,7 @@ private void assertTableRoute(String tableName, String logicalTableName, if (isOfflineExpected) { assertNotNull(logicalTableRouteInfo.getOfflineTables()); assertEquals(logicalTableRouteInfo.getOfflineTables().size(), 1); - Map offlineRoutingTable = + Map offlineRoutingTable = logicalTableRouteInfo.getOfflineTables().get(0).getOfflineRoutingTable(); assertNotNull(offlineRoutingTable); assertRoutingTableEqual(offlineRoutingTable, expectedOfflineRoutingTable); @@ -77,7 +77,7 @@ private void assertTableRoute(String tableName, String logicalTableName, if (isRealtimeExpected) { assertNotNull(logicalTableRouteInfo.getRealtimeTables()); assertEquals(logicalTableRouteInfo.getRealtimeTables().size(), 1); - Map realtimeRoutingTable = + Map realtimeRoutingTable = logicalTableRouteInfo.getRealtimeTables().get(0).getRealtimeRoutingTable(); assertNotNull(realtimeRoutingTable); assertRoutingTableEqual(realtimeRoutingTable, expectedRealtimeRoutingTable); diff --git a/pinot-core/src/test/java/org/apache/pinot/core/routing/MockRoutingManagerFactory.java b/pinot-core/src/test/java/org/apache/pinot/core/routing/MockRoutingManagerFactory.java index fba9f7945075..6a1e44f30aeb 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/routing/MockRoutingManagerFactory.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/routing/MockRoutingManagerFactory.java @@ -106,7 +106,7 @@ public RoutingManager buildRoutingManager( for (Map.Entry>> entry : _tableSegmentServersMap.entrySet()) { String tableNameWithType = entry.getKey(); Map> segmentServersMap = entry.getValue(); - Map serverRouteInfoMap = new HashMap<>(); + Map segmentsToQueryMap = new HashMap<>(); for (Map.Entry> segmentServersEntry : segmentServersMap.entrySet()) { String segment = segmentServersEntry.getKey(); List servers = segmentServersEntry.getValue(); @@ -118,11 +118,11 @@ public RoutingManager buildRoutingManager( server = servers.get(serverId % numServers); serverId++; } - serverRouteInfoMap.computeIfAbsent(server, k -> new ServerRouteInfo(new ArrayList<>(), null)) + segmentsToQueryMap.computeIfAbsent(server, k -> new SegmentsToQuery(new ArrayList<>(), null)) .getSegments() .add(segment); } - routingTableMap.put(tableNameWithType, new RoutingTable(serverRouteInfoMap, List.of(), 0)); + routingTableMap.put(tableNameWithType, new RoutingTable(segmentsToQueryMap, List.of(), 0)); tableSegmentsMap.put(tableNameWithType, new ArrayList<>(segmentServersMap.keySet())); } Map tablePartitionInfoMap = null; diff --git a/pinot-core/src/test/java/org/apache/pinot/core/transport/QueryRoutingTest.java b/pinot-core/src/test/java/org/apache/pinot/core/transport/QueryRoutingTest.java index 6bd4e35e10c3..f3ce7148b1d5 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/transport/QueryRoutingTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/transport/QueryRoutingTest.java @@ -32,7 +32,7 @@ import org.apache.pinot.core.common.datatable.DataTableBuilder; import org.apache.pinot.core.common.datatable.DataTableBuilderFactory; import org.apache.pinot.core.query.scheduler.QueryScheduler; -import org.apache.pinot.core.routing.ServerRouteInfo; +import org.apache.pinot.core.routing.SegmentsToQuery; import org.apache.pinot.core.transport.server.routing.stats.ServerRoutingStatsManager; import org.apache.pinot.server.access.AccessControl; import org.apache.pinot.spi.config.table.TableType; @@ -66,9 +66,9 @@ public class QueryRoutingTest { SERVER_INSTANCE.toServerRoutingInstance(TableType.REALTIME, ServerInstance.RoutingType.NETTY); private static final BrokerRequest BROKER_REQUEST = CalciteSqlCompiler.compileToBrokerRequest("SELECT * FROM testTable"); - private static final Map ROUTING_TABLE = + private static final Map ROUTING_TABLE = Collections.singletonMap(SERVER_INSTANCE, - new ServerRouteInfo(Collections.emptyList(), Collections.emptyList())); + new SegmentsToQuery(Collections.emptyList(), Collections.emptyList())); private QueryRouter _queryRouter; private ServerRoutingStatsManager _serverRoutingStatsManager; @@ -480,9 +480,9 @@ public void testSkipUnavailableServer() serverInstance1.toServerRoutingInstance(TableType.OFFLINE, ServerInstance.RoutingType.NETTY); ServerRoutingInstance serverRoutingInstance2 = serverInstance2.toServerRoutingInstance(TableType.OFFLINE, ServerInstance.RoutingType.NETTY); - Map routingTable = - Map.of(serverInstance1, new ServerRouteInfo(Collections.emptyList(), Collections.emptyList()), - serverInstance2, new ServerRouteInfo(Collections.emptyList(), Collections.emptyList())); + Map routingTable = + Map.of(serverInstance1, new SegmentsToQuery(Collections.emptyList(), Collections.emptyList()), + serverInstance2, new SegmentsToQuery(Collections.emptyList(), Collections.emptyList())); long requestId = 123; DataSchema dataSchema = diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LeafStageWorkerAssignmentRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LeafStageWorkerAssignmentRule.java index 3461a5388418..b4bd796c9002 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LeafStageWorkerAssignmentRule.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LeafStageWorkerAssignmentRule.java @@ -51,7 +51,7 @@ import org.apache.pinot.common.utils.config.QueryOptionsUtils; import org.apache.pinot.core.routing.RoutingManager; import org.apache.pinot.core.routing.RoutingTable; -import org.apache.pinot.core.routing.ServerRouteInfo; +import org.apache.pinot.core.routing.SegmentsToQuery; import org.apache.pinot.core.routing.TablePartitionInfo; import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; import org.apache.pinot.core.transport.ServerInstance; @@ -158,8 +158,8 @@ private PhysicalTableScan assignTableScan(PhysicalTableScan tableScan, long requ String tableType = routingEntry.getKey(); RoutingTable routingTable = routingEntry.getValue(); Map> currentSegmentsMap = new HashMap<>(); - Map tmp = routingTable.getServerInstanceToSegmentsMap(); - for (Map.Entry serverEntry : tmp.entrySet()) { + Map tmp = routingTable.getServerInstanceToSegmentsMap(); + for (Map.Entry serverEntry : tmp.entrySet()) { // TODO: Optional segments are not supported yet by the MSE. String instanceId = serverEntry.getKey().getInstanceId(); Preconditions.checkState(currentSegmentsMap.put(instanceId, serverEntry.getValue().getSegments()) == null, diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java index 00e08eef6de9..b4a77cad575e 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/routing/WorkerManager.java @@ -41,7 +41,7 @@ import org.apache.pinot.core.routing.LogicalTableRouteProvider; import org.apache.pinot.core.routing.RoutingManager; import org.apache.pinot.core.routing.RoutingTable; -import org.apache.pinot.core.routing.ServerRouteInfo; +import org.apache.pinot.core.routing.SegmentsToQuery; import org.apache.pinot.core.routing.TablePartitionReplicatedServersInfo; import org.apache.pinot.core.routing.TableRouteInfo; import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; @@ -472,8 +472,8 @@ private void assignWorkersToNonPartitionedLeafFragment(DispatchablePlanMetadata String tableType = routingEntry.getKey(); RoutingTable routingTable = routingEntry.getValue(); // for each server instance, attach all table types and their associated segment list. - Map segmentsMap = routingTable.getServerInstanceToSegmentsMap(); - for (Map.Entry serverEntry : segmentsMap.entrySet()) { + Map segmentsMap = routingTable.getServerInstanceToSegmentsMap(); + for (Map.Entry serverEntry : segmentsMap.entrySet()) { Map> tableTypeToSegmentListMap = serverInstanceToSegmentsMap.computeIfAbsent(serverEntry.getKey(), k -> new HashMap<>()); // TODO: support optional segments for multi-stage engine. @@ -656,9 +656,9 @@ private static void assignTableSegmentsToWorkers(LogicalTableRouteInfo logicalTa } private static void transferToServerInstanceLogicalSegmentsMap(String physicalTableName, - Map segmentsMap, + Map segmentsMap, Map>> serverInstanceToLogicalSegmentsMap) { - for (Map.Entry serverEntry : segmentsMap.entrySet()) { + for (Map.Entry serverEntry : segmentsMap.entrySet()) { Map> tableNameToSegmentsMap = serverInstanceToLogicalSegmentsMap.computeIfAbsent(serverEntry.getKey(), k -> new HashMap<>()); // TODO: support optional segments for multi-stage engine. From 4ace88c13fc41f39061e543659c3b6c6c797b4ef Mon Sep 17 00:00:00 2001 From: Gonzalo Ortiz Jaureguizar Date: Thu, 31 Jul 2025 16:10:32 +0200 Subject: [PATCH 053/167] mse: change SendStatsPredicate.Safe to only send when the cluster is homogeneous. (#16474) --- .../starter/helix/SendStatsPredicate.java | 29 +++---------------- 1 file changed, 4 insertions(+), 25 deletions(-) diff --git a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/SendStatsPredicate.java b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/SendStatsPredicate.java index 4678cab2dd6e..27c610a83e93 100644 --- a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/SendStatsPredicate.java +++ b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/SendStatsPredicate.java @@ -78,6 +78,7 @@ public static SendStatsPredicate create(PinotConfiguration serverConf, HelixMana } public enum Mode { + /// Sends stats only if all the cluster participants use the same known version. SAFE { @Override public SendStatsPredicate create(HelixManager helixManager) { @@ -210,6 +211,8 @@ public synchronized void onInstanceConfigChange(List instanceCon } else { LOGGER.warn("Send MSE stats is now disabled (problematic versions: {})", _problematicVersionsById); } + } else if (!_sendStats) { + LOGGER.info("Send MSE stats is still disabled (problematic versions: {})", _problematicVersionsById); } } @@ -223,10 +226,6 @@ private String getVersion(InstanceConfig instanceConfig) { return instanceConfig.getRecord().getStringField(CommonConstants.Helix.Instance.PINOT_VERSION_KEY, null); } - /// Returns true if the version is problematic - /// - /// Ideally [PinotVersion] should have a way to extract versions in comparable format, but given it doesn't we - /// need to parse the string here. In case version doesn't match `1\.x\..*`, we treat is as a problematic version private boolean isProblematicVersion(@Nullable String versionStr) { if (versionStr == null) { return true; @@ -234,27 +233,7 @@ private boolean isProblematicVersion(@Nullable String versionStr) { if (versionStr.equals(PinotVersion.UNKNOWN)) { return true; } - if (versionStr.equals(PinotVersion.VERSION)) { - return false; - } - // Lets try to parse 1.x versions - String[] splits = versionStr.trim().split("\\."); - if (splits.length < 2) { - return true; - } - // Versions less than 1.x are problematic for sure - if (!splits[0].equals("1")) { - return true; - } - try { - // Versions less than 1.4 are problematic - if (Integer.parseInt(splits[1]) < 4) { - return true; - } - } catch (NumberFormatException e) { - return true; - } - return false; + return !versionStr.equals(PinotVersion.VERSION); } } } From 132f90a0b40f8168513d279a21275279122fbfd9 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 31 Jul 2025 10:13:36 -0700 Subject: [PATCH 054/167] Bump software.amazon.awssdk:bom from 2.32.11 to 2.32.12 (#16472) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 59302ad4813c..96fb774b7fdb 100644 --- a/pom.xml +++ b/pom.xml @@ -178,7 +178,7 @@ 0.15.1 0.4.7 4.2.2 - 2.32.11 + 2.32.12 1.2.36 1.22.0 2.14.0 From d1926a6eda6ff0eef21d7a17ccf5a3d183f2ea12 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 31 Jul 2025 10:13:56 -0700 Subject: [PATCH 055/167] Bump org.apache.commons:commons-csv from 1.14.0 to 1.14.1 (#16473) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 96fb774b7fdb..6de6264fdb70 100644 --- a/pom.xml +++ b/pom.xml @@ -207,7 +207,7 @@ 1.14.0 1.28.0 3.6.1 - 1.14.0 + 1.14.1 2.12.0 1.11.0 2.20.0 From 62c75763944aeacfce1bd8eef7ce9a2087c5d612 Mon Sep 17 00:00:00 2001 From: avshenuk <4578760+avshenuk@users.noreply.github.com> Date: Thu, 31 Jul 2025 19:25:48 +0200 Subject: [PATCH 056/167] Fix unnesting of non-Map value collections. (#16387) --- .../ComplexTypeTransformer.java | 11 ++++- .../ComplexTypeTransformerTest.java | 45 +++++++++++++++++++ 2 files changed, 55 insertions(+), 1 deletion(-) diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/ComplexTypeTransformer.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/ComplexTypeTransformer.java index c19600cdb8af..6a94db718322 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/ComplexTypeTransformer.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/recordtransformer/ComplexTypeTransformer.java @@ -201,7 +201,12 @@ public List transform(List records) { for (String field : _fieldsToUnnest) { unnestedRecords = unnestCollection(unnestedRecords, field); } - unnestedRecords.forEach(unnestedRecord -> unnestedRecord.getFieldToValueMap().putAll(originalValues)); + unnestedRecords.forEach(unnestedRecord -> { + Map values = unnestedRecord.getFieldToValueMap(); + for (Map.Entry entry : originalValues.entrySet()) { + values.putIfAbsent(entry.getKey(), entry.getValue()); + } + }); if (record.isIncomplete()) { unnestedRecords.forEach(GenericRow::markIncomplete); } @@ -246,6 +251,8 @@ private void unnestCollection(GenericRow record, String column, List // use the record itself list.add(record); } else { + // Remove the value before flattening since we are going to add the flattened items + record.removeValue(column); for (Object obj : (Collection) value) { GenericRow copy = flattenCollectionItem(record, obj, column); list.add(copy); @@ -256,6 +263,8 @@ private void unnestCollection(GenericRow record, String column, List // use the record itself list.add(record); } else { + // Remove the value before flattening since we are going to add the flattened items + record.removeValue(column); for (Object obj : (Object[]) value) { GenericRow copy = flattenCollectionItem(record, obj, column); list.add(copy); diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/recordtransformer/ComplexTypeTransformerTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/recordtransformer/ComplexTypeTransformerTest.java index e3fab5f5e3ee..5183bd7bbbaa 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/recordtransformer/ComplexTypeTransformerTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/recordtransformer/ComplexTypeTransformerTest.java @@ -131,7 +131,9 @@ public void testUnnestCollection() { transformedRows = transformer.transform(List.of(genericRow)); assertEquals(transformedRows.size(), 2); assertEquals(transformedRows.get(0).getValue("array.a"), "v1"); + assertEquals(transformedRows.get(0).getValue("array"), array); assertEquals(transformedRows.get(1).getValue("array.a"), "v2"); + assertEquals(transformedRows.get(1).getValue("array"), array); // unnest sibling collections // { @@ -180,12 +182,20 @@ public void testUnnestCollection() { assertEquals(transformedRows.size(), 4); assertEquals(transformedRows.get(0).getValue("array.a"), "v1"); assertEquals(transformedRows.get(0).getValue("array2.b"), "v3"); + assertEquals(transformedRows.get(0).getValue("array"), array); + assertEquals(transformedRows.get(0).getValue("array2"), array2); assertEquals(transformedRows.get(1).getValue("array.a"), "v1"); assertEquals(transformedRows.get(1).getValue("array2.b"), "v4"); + assertEquals(transformedRows.get(1).getValue("array"), array); + assertEquals(transformedRows.get(1).getValue("array2"), array2); assertEquals(transformedRows.get(2).getValue("array.a"), "v2"); assertEquals(transformedRows.get(2).getValue("array2.b"), "v3"); + assertEquals(transformedRows.get(2).getValue("array"), array); + assertEquals(transformedRows.get(2).getValue("array2"), array2); assertEquals(transformedRows.get(3).getValue("array.a"), "v2"); assertEquals(transformedRows.get(3).getValue("array2.b"), "v4"); + assertEquals(transformedRows.get(3).getValue("array"), array); + assertEquals(transformedRows.get(3).getValue("array2"), array2); // unnest nested collection // { @@ -242,6 +252,41 @@ public void testUnnestCollection() { assertEquals(transformedRows.get(0).getValue("array.a"), "v1"); assertEquals(transformedRows.get(0).getValue("array.array2"), "[{\"b\":\"v3\"},{\"b\":\"v4\"}]"); assertEquals(transformedRows.get(1).getValue("array.a"), "v2"); + + // unnest root level collection with simple non-primitive values + // { + // "a": "value", + // "b": "another", + // "array":["x", "y"] + // } + // -> + // [{ + // "a": "value", + // "b": "another", + // "array":"x" + // }, + // { + // "a": "value", + // "b": "another", + // "array":"y" + // }] + genericRow = new GenericRow(); + genericRow.putValue("a", "value"); + genericRow.putValue("b", "another"); + array = new Object[2]; + array[0] = "x"; + array[1] = "y"; + genericRow.putValue("array", array); + transformer = + new ComplexTypeTransformer.Builder().setFieldsToUnnest(List.of("array")).build(); + transformedRows = transformer.transform(List.of(genericRow)); + assertEquals(transformedRows.size(), 2); + assertEquals(transformedRows.get(0).getValue("a"), "value"); + assertEquals(transformedRows.get(0).getValue("b"), "another"); + assertEquals(transformedRows.get(0).getValue("array"), "x"); + assertEquals(transformedRows.get(1).getValue("a"), "value"); + assertEquals(transformedRows.get(1).getValue("b"), "another"); + assertEquals(transformedRows.get(1).getValue("array"), "y"); } @Test From 1d534b41e3a4b26d0b4e39f154163074ef72822d Mon Sep 17 00:00:00 2001 From: RAGHVENDRA KUMAR YADAV Date: Thu, 31 Jul 2025 23:17:56 -0700 Subject: [PATCH 057/167] Adding Match prefix phrase query lucene parser (#16476) --- .../pinot/queries/TextSearchQueriesTest.java | 54 +++- .../parsers/PrefixPhraseQueryParser.java | 298 ++++++++++++++++++ .../local/utils/LuceneTextIndexUtils.java | 23 +- .../local/utils/LuceneTextIndexUtilsTest.java | 133 ++++++++ 4 files changed, 506 insertions(+), 2 deletions(-) create mode 100644 pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/text/lucene/parsers/PrefixPhraseQueryParser.java diff --git a/pinot-core/src/test/java/org/apache/pinot/queries/TextSearchQueriesTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/TextSearchQueriesTest.java index 0852bba44cae..b88276963a99 100644 --- a/pinot-core/src/test/java/org/apache/pinot/queries/TextSearchQueriesTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/queries/TextSearchQueriesTest.java @@ -2028,7 +2028,11 @@ public void testTextSearchWithOptions() }); String query = "SELECT INT_COL, SKILLS_TEXT_COL FROM " + TABLE_NAME + " WHERE TEXT_MATCH(" + SKILLS_TEXT_COL_NAME - + ", '*ealtime streaming system*', 'parser=CLASSIC,allowLeadingWildcard=true,defaultOperator=AND') LIMIT 50000"; + + ", 'realtime streaming system', 'parser=MATCHPHRASE') LIMIT 50000"; + testTextSearchSelectQueryHelper(query, 0, false, expected); + + query = "SELECT INT_COL, SKILLS_TEXT_COL FROM " + TABLE_NAME + " WHERE TEXT_MATCH(" + SKILLS_TEXT_COL_NAME + + ", 'realtime streaming system', 'parser=MATCHPHRASE,enablePrefixMatch=true') LIMIT 50000"; testTextSearchSelectQueryHelper(query, expected.size(), false, expected); List expected1 = new ArrayList<>(); @@ -2082,6 +2086,54 @@ public void testTextSearchWithOptions() testTextSearchSelectQueryHelper(query8, expected.size(), false, expected); } + @Test + public void testMatchPhraseQueryParser() + throws Exception { + // Test case 1: "Tensor flow" - should match 3 documents + List expectedTensorFlow = new ArrayList<>(); + expectedTensorFlow.add(new Object[]{ + 1004, "Machine learning, Tensor flow, Java, Stanford university," + }); + expectedTensorFlow.add(new Object[]{ + 1007, "C++, Python, Tensor flow, database kernel, storage, indexing and transaction processing, building " + + "large scale systems, Machine learning" + }); + expectedTensorFlow.add(new Object[]{ + 1016, "CUDA, GPU processing, Tensor flow, Pandas, Python, Jupyter notebook, spark, Machine learning, building" + + " high performance scalable systems" + }); + + // Test exact phrase "Tensor flow" with default settings (slop=0, inOrder=true) + String queryExactPhrase = + "SELECT INT_COL, SKILLS_TEXT_COL FROM " + TABLE_NAME + " WHERE TEXT_MATCH(" + SKILLS_TEXT_COL_NAME + + ", 'Tensor flow', 'parser=MATCHPHRASE,enablePrefixMatch=true') LIMIT 50000"; + testTextSearchSelectQueryHelper(queryExactPhrase, 3, false, expectedTensorFlow); + + // Test "Tensor database" with slop=1 (should allow one position gap) + List expectedTensorDatabase = new ArrayList<>(); + expectedTensorDatabase.add(new Object[]{ + 1007, "C++, Python, Tensor flow, database kernel, storage, indexing and transaction processing, building " + + "large scale systems, Machine learning" + }); + + String querySlop1 = + "SELECT INT_COL, SKILLS_TEXT_COL FROM " + TABLE_NAME + " WHERE TEXT_MATCH(" + SKILLS_TEXT_COL_NAME + + ", 'Tensor database', 'parser=MATCHPHRASE,enablePrefixMatch=true,slop=1') LIMIT 50000"; + testTextSearchSelectQueryHelper(querySlop1, 1, false, expectedTensorDatabase); + + // Test "Tensor flow" with inOrder=false (should allow any order) + String queryInOrderFalse = + "SELECT INT_COL, SKILLS_TEXT_COL FROM " + TABLE_NAME + " WHERE TEXT_MATCH(" + SKILLS_TEXT_COL_NAME + + ", 'Tensor flow', 'parser=MATCHPHRASE,enablePrefixMatch=true,inOrder=false') LIMIT 50000"; + testTextSearchSelectQueryHelper(queryInOrderFalse, 3, false, expectedTensorFlow); + + // Test "Tensor flow" with both slop=1 and inOrder=false + String querySlopAndInOrder = + "SELECT INT_COL, SKILLS_TEXT_COL FROM " + TABLE_NAME + " WHERE TEXT_MATCH(" + SKILLS_TEXT_COL_NAME + + ", 'flow Tensor', 'parser=MATCHPHRASE,enablePrefixMatch=true,inOrder=false') LIMIT 50000"; + testTextSearchSelectQueryHelper(querySlopAndInOrder, 3, false, expectedTensorFlow); + } + // ===== TEST CASES FOR AND/OR FILTER OPERATORS ===== @Test public void testTextSearchWithOptionsAndOrOperators() diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/text/lucene/parsers/PrefixPhraseQueryParser.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/text/lucene/parsers/PrefixPhraseQueryParser.java new file mode 100644 index 000000000000..d6c5b740433c --- /dev/null +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/text/lucene/parsers/PrefixPhraseQueryParser.java @@ -0,0 +1,298 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.segment.local.segment.index.text.lucene.parsers; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import org.apache.lucene.analysis.Analyzer; +import org.apache.lucene.analysis.TokenStream; +import org.apache.lucene.analysis.tokenattributes.CharTermAttribute; +import org.apache.lucene.index.Term; +import org.apache.lucene.queries.spans.SpanMultiTermQueryWrapper; +import org.apache.lucene.queries.spans.SpanNearQuery; +import org.apache.lucene.queries.spans.SpanQuery; +import org.apache.lucene.queries.spans.SpanTermQuery; +import org.apache.lucene.queryparser.charstream.CharStream; +import org.apache.lucene.queryparser.classic.ParseException; +import org.apache.lucene.queryparser.classic.QueryParserBase; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.WildcardQuery; + + +/** + * A custom query parser that creates prefix phrase queries. + * This parser tokenizes the input query and creates a SpanNearQuery where + * all terms except the last one are exact matches, and the last term can optionally + * have a wildcard suffix based on the enablePrefixMatch setting. + * + *

    This parser is designed to support both exact phrase matching and prefix phrase matching:

    + *
      + *
    • Exact phrase matching (default): All terms are matched exactly as they appear
      • + *
    • Prefix phrase matching: The last term is treated as a prefix with wildcard
      • + *
    + * + *

    Example usage:

    + *
      + *
    • Input: 'java realtime streaming' with enablePrefixMatch=false (default) + *
      Output: SpanNearQuery with exact matches for "java", "realtime", and "streaming"
      • + *
    • Input: 'java realtime streaming' with enablePrefixMatch=true + *
      Output: SpanNearQuery with exact matches for "java" and "realtime", + * and wildcard match for "streaming*"
      • + *
    • Input: 'stream' with enablePrefixMatch=false (default) + *
      Output: SpanTermQuery for exact match "stream"
      • + *
    • Input: 'stream' with enablePrefixMatch=true + *
      Output: SpanMultiTermQueryWrapper for wildcard match "stream*"
      • + *
    + * + *

    Behavior:

    + *
      + *
    • Single term queries: Returns SpanTermQuery (exact) or SpanMultiTermQueryWrapper (prefix)
      • + *
    • Multiple term queries: Returns SpanNearQuery with all terms in exact order
      • + *
    • Null/empty queries: Throws ParseException
      • + *
    • Whitespace-only queries: Throws ParseException
      • + *
    + * + *

    This parser extends Lucene's QueryParserBase and implements the required abstract methods. + * It uses the provided Analyzer for tokenization and creates appropriate Lucene Span queries.

    + */ +public class PrefixPhraseQueryParser extends QueryParserBase { + /** The field name to search in */ + private final String _field; + + /** The analyzer used for tokenizing the query */ + private final Analyzer _analyzer; + + /** Flag to control whether prefix matching is enabled on the last term */ + private boolean _enablePrefixMatch = false; + + /** The slop (distance) allowed between terms in the phrase query. Default is 0 (exact order) */ + private int _slop = 0; + + /** Whether terms must appear in the specified order. Default is true (exact order) */ + private boolean _inOrder = true; + + /** + * Constructs a new PrefixPhraseQueryParser with the specified field and analyzer. + * + * @param field the field name to search in (must not be null) + * @param analyzer the analyzer to use for tokenizing queries (must not be null) + * @throws IllegalArgumentException if field or analyzer is null + */ + public PrefixPhraseQueryParser(String field, Analyzer analyzer) { + super(); + _field = field; + _analyzer = analyzer; + } + + /** + * Sets whether to enable prefix matching on the last term. + * + *

    When enabled ({@code true}): + *

      + *
    • Single term queries: Returns a SpanMultiTermQueryWrapper with wildcard (*)
      • + *
    • Multiple term queries: The last term gets a wildcard suffix (*)
      • + *
    + * + *

    When disabled ({@code false}, default): + *

      + *
    • Single term queries: Returns a SpanTermQuery for exact match
      • + *
    • Multiple term queries: All terms are matched exactly
      • + *
    + * + * @param enablePrefixMatch true to enable prefix matching, false to disable (default) + */ + public void setEnablePrefixMatch(boolean enablePrefixMatch) { + _enablePrefixMatch = enablePrefixMatch; + } + + /** + * Sets the slop (distance) allowed between terms in the phrase query. + * + *

    The slop determines how many positions apart the terms can be while still matching. + * For example:

    + *
      + *
    • slop=0: Terms must be adjacent in exact order
      • + *
    • slop=1: Terms can be 1 position apart
      • + *
    • slop=2: Terms can be 2 positions apart
      • + *
    + * + *

    This setting only affects multiple term queries that create SpanNearQuery.

    + * + * @param slop the number of positions allowed between terms (default is 0) + * @throws IllegalArgumentException if slop is negative + */ + public void setSlop(int slop) { + if (slop < 0) { + throw new IllegalArgumentException("Slop cannot be negative: " + slop); + } + _slop = slop; + } + + /** + * Sets whether terms must appear in the specified order. + * + *

    When enabled ({@code true}, default): + *

      + *
    • Terms must appear in the exact order specified in the query
      • + *
    • Example: "java realtime" matches "java realtime streaming" but not "realtime java streaming"
      • + *
    + * + *

    When disabled ({@code false}): + *

      + *
    • Terms can appear in any order within the slop distance
      • + *
    • Example: "java realtime" matches both "java realtime streaming" and "realtime java streaming"
      • + *
    + * + *

    This setting only affects multiple term queries that create SpanNearQuery.

    + * + * @param inOrder true to require terms in exact order, false to allow any order + */ + public void setInOrder(boolean inOrder) { + _inOrder = inOrder; + } + + /** + * Parses the given query string and returns an appropriate Lucene Query. + * + *

    This method performs the following steps:

    + *
      + *
    1. Validates the input query (null, empty, whitespace-only)
      2. + *
    2. Tokenizes the query using the configured analyzer
      3. + *
    3. Creates appropriate Lucene queries based on the number of tokens and enablePrefixMatch setting
      4. + *
    + * + *

    Query Types Returned:

    + *
      + *
    • Single term: + *
        + *
      • If enablePrefixMatch=false: SpanTermQuery for exact match
        • + *
      • If enablePrefixMatch=true: SpanMultiTermQueryWrapper with wildcard
        • + *
      + *
      • + *
    • Multiple terms: SpanNearQuery with all terms in exact order + *
        + *
      • All terms except the last: SpanTermQuery (exact match)
        • + *
      • Last term: SpanTermQuery (exact) or SpanMultiTermQueryWrapper (wildcard) + * based on enablePrefixMatch
        • + *
      + *
      • + *
    + * + * @param query the query string to parse (must not be null or empty) + * @return a Lucene Query object representing the parsed query + * @throws ParseException if the query is null, empty, or contains no valid tokens after tokenization + * @throws RuntimeException if tokenization fails due to an IOException + */ + @Override + public Query parse(String query) throws ParseException { + if (query == null) { + throw new ParseException("Query cannot be null"); + } + + if (query.trim().isEmpty()) { + throw new ParseException("Query cannot be empty"); + } + + // Tokenize the query + List tokens = new ArrayList<>(); + try (TokenStream stream = _analyzer.tokenStream(_field, query)) { + stream.reset(); + CharTermAttribute charTermAttribute = stream.addAttribute(CharTermAttribute.class); + + while (stream.incrementToken()) { + String token = charTermAttribute.toString(); + if (!token.trim().isEmpty()) { + tokens.add(token); + } + } + stream.end(); + } catch (IOException e) { + throw new RuntimeException("Failed to tokenize query: " + query, e); + } + + // Check if we have any valid tokens after tokenization + if (tokens.isEmpty()) { + throw new ParseException("Query tokenization resulted in no valid tokens"); + } + + // Handle single token case + if (tokens.size() == 1) { + String token = tokens.get(0); + if (_enablePrefixMatch) { + WildcardQuery wildcardQuery = new WildcardQuery(new Term(_field, token + "*")); + return new SpanMultiTermQueryWrapper<>(wildcardQuery); + } else { + return new SpanTermQuery(new Term(_field, token)); + } + } + + // Handle multiple tokens case + List spanQueries = new ArrayList<>(); + + // Add regular SpanTermQueries for all tokens except the last one + for (int i = 0; i < tokens.size() - 1; i++) { + spanQueries.add(new SpanTermQuery(new Term(_field, tokens.get(i)))); + } + + // Add query for the last token + String lastToken = tokens.get(tokens.size() - 1); + if (_enablePrefixMatch) { + WildcardQuery wildcardQuery = new WildcardQuery(new Term(_field, lastToken + "*")); + spanQueries.add(new SpanMultiTermQueryWrapper<>(wildcardQuery)); + } else { + spanQueries.add(new SpanTermQuery(new Term(_field, lastToken))); + } + + // Create SpanNearQuery with configurable slop and inOrder settings + return new SpanNearQuery(spanQueries.toArray(new SpanQuery[0]), _slop, _inOrder); + } + + /** + * Reinitializes the parser with a new CharStream. + * + *

    This method is required by QueryParserBase but is not used in this implementation + * since we override the parse(String) method directly. The method is left as a no-op.

    + * + * @param input the CharStream to reinitialize with (ignored in this implementation) + */ + @Override + public void ReInit(CharStream input) { + // This method is required by QueryParserBase but not used in our implementation + // since we override parse(String) directly + } + + /** + * Creates a top-level query for the specified field. + * + *

    This method is required by QueryParserBase but is not supported in this implementation. + * Use the parse(String) method instead for query parsing.

    + * + * @param field the field name (ignored in this implementation) + * @return never returns (always throws UnsupportedOperationException) + * @throws ParseException never thrown (method always throws UnsupportedOperationException) + * @throws UnsupportedOperationException always thrown, indicating this method is not supported + */ + @Override + public Query TopLevelQuery(String field) + throws ParseException { + throw new UnsupportedOperationException( + "TopLevelQuery is not supported in PrefixPhraseQueryParser. Use parse(String) method instead."); + } +} diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/LuceneTextIndexUtils.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/LuceneTextIndexUtils.java index dfa4c9f6fb91..d366789d9c7f 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/LuceneTextIndexUtils.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/LuceneTextIndexUtils.java @@ -48,6 +48,7 @@ public class LuceneTextIndexUtils { public static final String PARSER_CLASSIC = "CLASSIC"; public static final String PARSER_STANDARD = "STANDARD"; public static final String PARSER_COMPLEX = "COMPLEX"; + public static final String PARSER_MATCHPHRASE = "MATCHPHRASE"; // Default operator constants public static final String DEFAULT_OPERATOR_AND = "AND"; @@ -76,6 +77,9 @@ public static class OptionKey { public static final String TIME_ZONE = "timeZone"; public static final String PHRASE_SLOP = "phraseSlop"; public static final String MAX_DETERMINIZED_STATES = "maxDeterminizedStates"; + public static final String SLOP = "slop"; + public static final String IN_ORDER = "inOrder"; + public static final String ENABLE_PREFIX_MATCH = "enablePrefixMatch"; } // Parser class names @@ -84,6 +88,8 @@ public static class OptionKey { public static final String COMPLEX_PHRASE_QUERY_PARSER_CLASS = "org.apache.lucene.queryparser.complexPhrase.ComplexPhraseQueryParser"; public static final String CLASSIC_QUERY_PARSER = "org.apache.lucene.queryparser.classic.QueryParser"; + public static final String MATCHPHRASE_QUERY_PARSER_CLASS = + "org.apache.pinot.segment.local.segment.index.text.lucene.parsers.PrefixPhraseQueryParser"; private LuceneTextIndexUtils() { } @@ -147,6 +153,9 @@ public static Query createQueryParserWithOptions(String actualQuery, LuceneTextI case PARSER_COMPLEX: parserClassName = COMPLEX_PHRASE_QUERY_PARSER_CLASS; break; + case PARSER_MATCHPHRASE: + parserClassName = MATCHPHRASE_QUERY_PARSER_CLASS; + break; default: parserClassName = CLASSIC_QUERY_PARSER; break; @@ -224,7 +233,7 @@ public static Query createQueryParserWithOptions(String actualQuery, LuceneTextI Method parseMethod = parser.getClass().getMethod("parse", String.class, String.class); query = (Query) parseMethod.invoke(parser, actualQuery, column); } else { - // Other parsers use parse(String) + // Other parsers (CLASSIC, COMPLEX, MATCHPHRASE) use parse(String) Method parseMethod = parser.getClass().getMethod("parse", String.class); query = (Query) parseMethod.invoke(parser, actualQuery); } @@ -332,6 +341,18 @@ public int getPhraseSlop() { public int getMaxDeterminizedStates() { return Integer.parseInt(_options.getOrDefault(OptionKey.MAX_DETERMINIZED_STATES, "10000")); } + + public int getSlop() { + return Integer.parseInt(_options.getOrDefault(OptionKey.SLOP, "0")); + } + + public boolean isInOrder() { + return Boolean.parseBoolean(_options.getOrDefault(OptionKey.IN_ORDER, "true")); + } + + public boolean isEnablePrefixMatch() { + return Boolean.parseBoolean(_options.getOrDefault(OptionKey.ENABLE_PREFIX_MATCH, "false")); + } } /** diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/LuceneTextIndexUtilsTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/LuceneTextIndexUtilsTest.java index 2308809584c4..b8a70caad670 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/LuceneTextIndexUtilsTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/LuceneTextIndexUtilsTest.java @@ -18,7 +18,10 @@ */ package org.apache.pinot.segment.local.utils; +import java.lang.reflect.InvocationTargetException; import java.util.Map; +import org.apache.lucene.analysis.Analyzer; +import org.apache.lucene.analysis.core.WhitespaceAnalyzer; import org.apache.lucene.index.Term; import org.apache.lucene.queries.spans.SpanMultiTermQueryWrapper; import org.apache.lucene.queries.spans.SpanNearQuery; @@ -27,9 +30,11 @@ import org.apache.lucene.search.BooleanClause; import org.apache.lucene.search.BooleanQuery; import org.apache.lucene.search.PrefixQuery; +import org.apache.lucene.search.Query; import org.apache.lucene.search.RegexpQuery; import org.apache.lucene.search.TermQuery; import org.apache.lucene.search.WildcardQuery; +import org.apache.pinot.segment.local.segment.index.text.lucene.parsers.PrefixPhraseQueryParser; import org.testng.Assert; import org.testng.annotations.Test; @@ -202,4 +207,132 @@ public void testLuceneTextIndexOptionsAllGetters() { Assert.assertEquals(options.getPhraseSlop(), 2); Assert.assertEquals(options.getMaxDeterminizedStates(), 5000); } + + @Test + public void testMatchPhraseQueryParser() + throws Exception { + // Test the new MATCHPHRASE parser functionality + String optionsString = "parser=MATCHPHRASE,enablePrefixMatch=true"; + LuceneTextIndexUtils.LuceneTextIndexOptions options = + new LuceneTextIndexUtils.LuceneTextIndexOptions(optionsString); + + // Create a simple analyzer for testing + Analyzer analyzer = new WhitespaceAnalyzer(); + String column = "testColumn"; + + // Test positive case: "java realtime streaming" + String query = "java realtime streaming"; + + Query result = LuceneTextIndexUtils.createQueryParserWithOptions(query, options, column, analyzer); + Assert.assertNotNull(result); + Assert.assertTrue(result instanceof SpanNearQuery); + + // Test positive case: "realtime stream*" + query = "realtime stream*"; + result = LuceneTextIndexUtils.createQueryParserWithOptions(query, options, column, analyzer); + Assert.assertNotNull(result); + Assert.assertTrue(result instanceof SpanNearQuery); + + // Test positive case: "stream*" - single term should return SpanMultiTermQueryWrapper + query = "stream*"; + result = LuceneTextIndexUtils.createQueryParserWithOptions(query, options, column, analyzer); + Assert.assertNotNull(result); + Assert.assertTrue(result instanceof SpanMultiTermQueryWrapper); + + // Test edge case: empty string "" + query = ""; + try { + LuceneTextIndexUtils.createQueryParserWithOptions(query, options, column, analyzer); + Assert.fail("Expected exception for empty query"); + } catch (RuntimeException e) { + // The method wraps ParseException in RuntimeException via reflection + Assert.assertTrue(e.getCause() instanceof InvocationTargetException); + } + + // Test edge case: null query + try { + LuceneTextIndexUtils.createQueryParserWithOptions(null, options, column, analyzer); + Assert.fail("Expected exception for null query"); + } catch (RuntimeException e) { + // The method wraps ParseException in RuntimeException via reflection + Assert.assertTrue(e.getCause() instanceof InvocationTargetException); + } + + // Test that TopLevelQuery throws UnsupportedOperationException + try { + PrefixPhraseQueryParser parser = new PrefixPhraseQueryParser(column, analyzer); + parser.TopLevelQuery(column); + Assert.fail("Expected UnsupportedOperationException for TopLevelQuery"); + } catch (UnsupportedOperationException e) { + Assert.assertTrue(e.getMessage().contains("TopLevelQuery is not supported")); + } + + // Test slop and inOrder settings + PrefixPhraseQueryParser slopParser = new PrefixPhraseQueryParser(column, analyzer); + + // Test default slop and inOrder (0 slop, true inOrder) + Query defaultSlopQuery = slopParser.parse("java realtime streaming"); + Assert.assertTrue(defaultSlopQuery instanceof SpanNearQuery); + + // Test custom slop and inOrder + slopParser.setSlop(2); + slopParser.setInOrder(false); + Query customSlopQuery = slopParser.parse("java realtime streaming"); + Assert.assertTrue(customSlopQuery instanceof SpanNearQuery); + + // Test invalid slop (should throw exception) + try { + slopParser.setSlop(-1); + Assert.fail("Expected IllegalArgumentException for negative slop"); + } catch (IllegalArgumentException e) { + Assert.assertTrue(e.getMessage().contains("Slop cannot be negative")); + } + + // Test slop and inOrder with createQueryParserWithOptions + LuceneTextIndexUtils.LuceneTextIndexOptions slopOptions = + LuceneTextIndexUtils.createOptions("parser=MATCHPHRASE,enablePrefixMatch=true"); + + // Test default slop and inOrder behavior + Query defaultSlopResult = LuceneTextIndexUtils.createQueryParserWithOptions( + "java realtime streaming", slopOptions, column, analyzer); + Assert.assertTrue(defaultSlopResult instanceof SpanNearQuery); + + // Test custom slop and inOrder settings + LuceneTextIndexUtils.LuceneTextIndexOptions customSlopOptions = + LuceneTextIndexUtils.createOptions("parser=MATCHPHRASE,enablePrefixMatch=true"); + + // Create a parser instance to test slop and inOrder settings + PrefixPhraseQueryParser customParser = new PrefixPhraseQueryParser(column, analyzer); + customParser.setEnablePrefixMatch(true); + customParser.setSlop(2); + customParser.setInOrder(false); + + // Test that custom settings work correctly + Query customSlopResult = customParser.parse("java realtime streaming"); + Assert.assertTrue(customSlopResult instanceof SpanNearQuery); + + // Test that the parser can be configured with different slop values + customParser.setSlop(1); + Query slop1Result = customParser.parse("java realtime streaming"); + Assert.assertTrue(slop1Result instanceof SpanNearQuery); + + // Test that the parser can be configured with different inOrder values + customParser.setInOrder(true); + Query inOrderTrueResult = customParser.parse("java realtime streaming"); + Assert.assertTrue(inOrderTrueResult instanceof SpanNearQuery); + + // Test default behavior using createOptions + LuceneTextIndexUtils.LuceneTextIndexOptions defaultOptions = + LuceneTextIndexUtils.createOptions("parser=MATCHPHRASE"); + + // Test single term with default behavior (prefix match disabled) + Query defaultSingleTermQuery = + LuceneTextIndexUtils.createQueryParserWithOptions("stream", defaultOptions, column, analyzer); + Assert.assertTrue(defaultSingleTermQuery instanceof SpanTermQuery); + + // Test multiple terms with default behavior (prefix match disabled) + Query defaultMultiTermQuery = + LuceneTextIndexUtils.createQueryParserWithOptions("java realtime streaming", defaultOptions, column, analyzer); + Assert.assertTrue(defaultMultiTermQuery instanceof SpanNearQuery); + } } From 2cf81becffbb1e4df5985f973219b451617f4646 Mon Sep 17 00:00:00 2001 From: "Xiaotian (Jackie) Jiang" <17555551+Jackie-Jiang@users.noreply.github.com> Date: Fri, 1 Aug 2025 00:35:04 -0600 Subject: [PATCH 058/167] Enhance REGEXP_LIKE to scan dictionary when dictionary is small (#16478) --- .../operator/filter/FilterOperatorUtils.java | 36 +++------- ...ctIdBasedRegexpLikePredicateEvaluator.java | 31 ++++++++ ...TBasedRegexpPredicateEvaluatorFactory.java | 2 +- ...TBasedRegexpPredicateEvaluatorFactory.java | 2 +- .../RegexpLikePredicateEvaluatorFactory.java | 70 +++++++++++++++++-- .../tests/OfflineClusterIntegrationTest.java | 19 +++++ 6 files changed, 127 insertions(+), 33 deletions(-) create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/BaseDictIdBasedRegexpLikePredicateEvaluator.java diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/FilterOperatorUtils.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/FilterOperatorUtils.java index b6cab14d83dd..2f38d1844315 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/FilterOperatorUtils.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/FilterOperatorUtils.java @@ -23,7 +23,7 @@ import java.util.List; import java.util.OptionalInt; import org.apache.pinot.common.request.context.predicate.Predicate; -import org.apache.pinot.common.request.context.predicate.RegexpLikePredicate; +import org.apache.pinot.core.operator.filter.predicate.BaseDictIdBasedRegexpLikePredicateEvaluator; import org.apache.pinot.core.operator.filter.predicate.PredicateEvaluator; import org.apache.pinot.core.query.request.context.QueryContext; import org.apache.pinot.segment.spi.datasource.DataSource; @@ -54,20 +54,19 @@ BaseFilterOperator getLeafFilterOperator(QueryContext queryContext, PredicateEva /** * Returns the AND filter operator or equivalent filter operator. */ - BaseFilterOperator getAndFilterOperator(QueryContext queryContext, - List filterOperators, int numDocs); + BaseFilterOperator getAndFilterOperator(QueryContext queryContext, List filterOperators, + int numDocs); /** * Returns the OR filter operator or equivalent filter operator. */ - BaseFilterOperator getOrFilterOperator(QueryContext queryContext, - List filterOperators, int numDocs); + BaseFilterOperator getOrFilterOperator(QueryContext queryContext, List filterOperators, + int numDocs); /** * Returns the NOT filter operator or equivalent filter operator. */ - BaseFilterOperator getNotFilterOperator(QueryContext queryContext, BaseFilterOperator filterOperator, - int numDocs); + BaseFilterOperator getNotFilterOperator(QueryContext queryContext, BaseFilterOperator filterOperator, int numDocs); } public static class DefaultImplementation implements Implementation { @@ -107,24 +106,12 @@ public BaseFilterOperator getLeafFilterOperator(QueryContext queryContext, Predi } return new ScanBasedFilterOperator(queryContext, predicateEvaluator, dataSource, numDocs); } else if (predicateType == Predicate.Type.REGEXP_LIKE) { - // Check if case-insensitive flag is present - RegexpLikePredicate regexpLikePredicate = (RegexpLikePredicate) predicateEvaluator.getPredicate(); - boolean caseInsensitive = regexpLikePredicate.isCaseInsensitive(); - if (caseInsensitive) { - if (dataSource.getIFSTIndex() != null && dataSource.getDataSourceMetadata().isSorted() + if (predicateEvaluator instanceof BaseDictIdBasedRegexpLikePredicateEvaluator) { + if (dataSource.getDataSourceMetadata().isSorted() && queryContext.isIndexUseAllowed(dataSource, FieldConfig.IndexType.SORTED)) { return new SortedIndexBasedFilterOperator(queryContext, predicateEvaluator, dataSource, numDocs); } - if (dataSource.getIFSTIndex() != null && dataSource.getInvertedIndex() != null - && queryContext.isIndexUseAllowed(dataSource, FieldConfig.IndexType.INVERTED)) { - return new InvertedIndexFilterOperator(queryContext, predicateEvaluator, dataSource, numDocs); - } - } else { - if (dataSource.getFSTIndex() != null && dataSource.getDataSourceMetadata().isSorted() - && queryContext.isIndexUseAllowed(dataSource, FieldConfig.IndexType.SORTED)) { - return new SortedIndexBasedFilterOperator(queryContext, predicateEvaluator, dataSource, numDocs); - } - if (dataSource.getFSTIndex() != null && dataSource.getInvertedIndex() != null + if (dataSource.getInvertedIndex() != null && queryContext.isIndexUseAllowed(dataSource, FieldConfig.IndexType.INVERTED)) { return new InvertedIndexFilterOperator(queryContext, predicateEvaluator, dataSource, numDocs); } @@ -174,8 +161,8 @@ public BaseFilterOperator getAndFilterOperator(QueryContext queryContext, List filterOperators, int numDocs) { + public BaseFilterOperator getOrFilterOperator(QueryContext queryContext, List filterOperators, + int numDocs) { List childFilterOperators = new ArrayList<>(filterOperators.size()); for (BaseFilterOperator filterOperator : filterOperators) { if (filterOperator.isResultMatchingAll()) { @@ -210,7 +197,6 @@ public BaseFilterOperator getNotFilterOperator(QueryContext queryContext, BaseFi return new NotFilterOperator(filterOperator, numDocs, queryContext.isNullHandlingEnabled()); } - /** * For AND filter operator, reorders its child filter operators based on their cost and puts the ones with * inverted index first in order to reduce the number of documents to be processed. diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/BaseDictIdBasedRegexpLikePredicateEvaluator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/BaseDictIdBasedRegexpLikePredicateEvaluator.java new file mode 100644 index 000000000000..806737763570 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/BaseDictIdBasedRegexpLikePredicateEvaluator.java @@ -0,0 +1,31 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.core.operator.filter.predicate; + +import org.apache.pinot.common.request.context.predicate.Predicate; +import org.apache.pinot.segment.spi.index.reader.Dictionary; + + +/// Base class for dictionary-based REGEXP_LIKE predicate evaluators that use dictionary IDs for matching. +public abstract class BaseDictIdBasedRegexpLikePredicateEvaluator extends BaseDictionaryBasedPredicateEvaluator { + + protected BaseDictIdBasedRegexpLikePredicateEvaluator(Predicate predicate, Dictionary dictionary) { + super(predicate, dictionary); + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/FSTBasedRegexpPredicateEvaluatorFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/FSTBasedRegexpPredicateEvaluatorFactory.java index 11dbe7aa995c..8cbdb4288450 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/FSTBasedRegexpPredicateEvaluatorFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/FSTBasedRegexpPredicateEvaluatorFactory.java @@ -49,7 +49,7 @@ public static BaseDictionaryBasedPredicateEvaluator newFSTBasedEvaluator(RegexpL /** * Matches regexp query using FSTIndexReader. */ - private static class FSTBasedRegexpPredicateEvaluator extends BaseDictionaryBasedPredicateEvaluator { + private static class FSTBasedRegexpPredicateEvaluator extends BaseDictIdBasedRegexpLikePredicateEvaluator { final ImmutableRoaringBitmap _matchingDictIdBitmap; public FSTBasedRegexpPredicateEvaluator(RegexpLikePredicate regexpLikePredicate, TextIndexReader fstIndexReader, diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/IFSTBasedRegexpPredicateEvaluatorFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/IFSTBasedRegexpPredicateEvaluatorFactory.java index a2d82b765d3d..074d44d3e9b4 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/IFSTBasedRegexpPredicateEvaluatorFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/IFSTBasedRegexpPredicateEvaluatorFactory.java @@ -45,7 +45,7 @@ public static BaseDictionaryBasedPredicateEvaluator newIFSTBasedEvaluator(Regexp return new IFSTBasedRegexpPredicateEvaluator(regexpLikePredicate, ifstIndexReader, dictionary); } - private static class IFSTBasedRegexpPredicateEvaluator extends BaseDictionaryBasedPredicateEvaluator { + private static class IFSTBasedRegexpPredicateEvaluator extends BaseDictIdBasedRegexpLikePredicateEvaluator { final ImmutableRoaringBitmap _matchingDictIdBitmap; public IFSTBasedRegexpPredicateEvaluator(RegexpLikePredicate regexpLikePredicate, diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/RegexpLikePredicateEvaluatorFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/RegexpLikePredicateEvaluatorFactory.java index e9c2bb73fd95..169230831cff 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/RegexpLikePredicateEvaluatorFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/filter/predicate/RegexpLikePredicateEvaluatorFactory.java @@ -19,6 +19,10 @@ package org.apache.pinot.core.operator.filter.predicate; import com.google.common.base.Preconditions; +import it.unimi.dsi.fastutil.ints.IntArrayList; +import it.unimi.dsi.fastutil.ints.IntList; +import it.unimi.dsi.fastutil.ints.IntOpenHashSet; +import it.unimi.dsi.fastutil.ints.IntSet; import org.apache.pinot.common.request.context.predicate.RegexpLikePredicate; import org.apache.pinot.common.utils.regex.Matcher; import org.apache.pinot.segment.spi.index.reader.Dictionary; @@ -32,6 +36,9 @@ public class RegexpLikePredicateEvaluatorFactory { private RegexpLikePredicateEvaluatorFactory() { } + /// When the cardinality of the dictionary is less than this threshold, scan the dictionary to get the matching ids. + public static final int DICTIONARY_CARDINALITY_THRESHOLD_FOR_SCAN = 10000; + /** * Create a new instance of dictionary based REGEXP_LIKE predicate evaluator. * @@ -42,9 +49,12 @@ private RegexpLikePredicateEvaluatorFactory() { */ public static BaseDictionaryBasedPredicateEvaluator newDictionaryBasedEvaluator( RegexpLikePredicate regexpLikePredicate, Dictionary dictionary, DataType dataType) { - boolean condition = (dataType == DataType.STRING || dataType == DataType.JSON); - Preconditions.checkArgument(condition, "Unsupported data type: " + dataType); - return new DictionaryBasedRegexpLikePredicateEvaluator(regexpLikePredicate, dictionary); + Preconditions.checkArgument(dataType.getStoredType() == DataType.STRING, "Unsupported data type: " + dataType); + if (dictionary.length() < DICTIONARY_CARDINALITY_THRESHOLD_FOR_SCAN) { + return new DictIdBasedRegexpLikePredicateEvaluator(regexpLikePredicate, dictionary); + } else { + return new ScanBasedRegexpLikePredicateEvaluator(regexpLikePredicate, dictionary); + } } /** @@ -56,16 +66,64 @@ public static BaseDictionaryBasedPredicateEvaluator newDictionaryBasedEvaluator( */ public static BaseRawValueBasedPredicateEvaluator newRawValueBasedEvaluator(RegexpLikePredicate regexpLikePredicate, DataType dataType) { - Preconditions.checkArgument(dataType == DataType.STRING, "Unsupported data type: " + dataType); + Preconditions.checkArgument(dataType.getStoredType() == DataType.STRING, "Unsupported data type: " + dataType); return new RawValueBasedRegexpLikePredicateEvaluator(regexpLikePredicate); } - private static final class DictionaryBasedRegexpLikePredicateEvaluator extends BaseDictionaryBasedPredicateEvaluator { + private static final class DictIdBasedRegexpLikePredicateEvaluator + extends BaseDictIdBasedRegexpLikePredicateEvaluator { + final IntSet _matchingDictIdSet; + + public DictIdBasedRegexpLikePredicateEvaluator(RegexpLikePredicate regexpLikePredicate, Dictionary dictionary) { + super(regexpLikePredicate, dictionary); + Matcher matcher = regexpLikePredicate.getPattern().matcher(""); + IntList matchingDictIds = new IntArrayList(); + int dictionarySize = _dictionary.length(); + for (int dictId = 0; dictId < dictionarySize; dictId++) { + if (matcher.reset(dictionary.getStringValue(dictId)).find()) { + matchingDictIds.add(dictId); + } + } + int numMatchingDictIds = matchingDictIds.size(); + if (numMatchingDictIds == 0) { + _alwaysFalse = true; + } else if (dictionarySize == numMatchingDictIds) { + _alwaysTrue = true; + } + _matchingDictIds = matchingDictIds.toIntArray(); + _matchingDictIdSet = new IntOpenHashSet(_matchingDictIds); + } + + @Override + public int getNumMatchingItems() { + return _matchingDictIdSet.size(); + } + + @Override + public boolean applySV(int dictId) { + return _matchingDictIdSet.contains(dictId); + } + + @Override + public int applySV(int limit, int[] docIds, int[] values) { + // reimplemented here to ensure applySV can be inlined + int matches = 0; + for (int i = 0; i < limit; i++) { + int value = values[i]; + if (applySV(value)) { + docIds[matches++] = docIds[i]; + } + } + return matches; + } + } + + private static final class ScanBasedRegexpLikePredicateEvaluator extends BaseDictionaryBasedPredicateEvaluator { // Reuse matcher to avoid excessive allocation. This is safe to do because the evaluator is always used // within the scope of a single thread. final Matcher _matcher; - public DictionaryBasedRegexpLikePredicateEvaluator(RegexpLikePredicate regexpLikePredicate, Dictionary dictionary) { + public ScanBasedRegexpLikePredicateEvaluator(RegexpLikePredicate regexpLikePredicate, Dictionary dictionary) { super(regexpLikePredicate, dictionary); _matcher = regexpLikePredicate.getPattern().matcher(""); } diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterIntegrationTest.java index 36af42448233..83ea21464363 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterIntegrationTest.java @@ -1893,6 +1893,8 @@ public void testDefaultColumns(boolean useMultiStageQueryEngine) }, 60_000L, "Failed to query new added columns without reload"); // Table size shouldn't change without reload assertEquals(getTableSize(getTableName()), _tableSize); + // Test REGEXP_LIKE on new added columns + testRegexpLikeOnNewAddedColumns(); // Trigger reload and verify column count reloadAllSegments(TEST_EXTRA_COLUMNS_QUERY, false, numTotalDocs); @@ -1942,6 +1944,7 @@ public void testDefaultColumns(boolean useMultiStageQueryEngine) assertTrue(derivedNullStringColumnIndex.has(StandardIndexes.NULL_VALUE_VECTOR_ID)); testNewAddedColumns(); + testRegexpLikeOnNewAddedColumns(); // The multi-stage query engine doesn't support expression overrides currently if (!useMultiStageQueryEngine()) { @@ -2209,6 +2212,22 @@ private void testNewAddedColumns() assertEquals(row.get(12).asDouble(), 0.0); } + private void testRegexpLikeOnNewAddedColumns() + throws Exception { + int numTotalDocs = (int) getCountStarResult(); + + // REGEXP_LIKE on new added dictionary-encoded columns should not scan the table when it matches all or nothing + for (String column : List.of("NewAddedSVJSONDimension", "NewAddedDerivedNullString")) { + JsonNode response = postQuery("SELECT COUNT(*) FROM mytable WHERE REGEXP_LIKE(" + column + ", 'foo')"); + assertEquals(response.get("resultTable").get("rows").get(0).get(0).asInt(), 0); + assertEquals(response.get("numEntriesScannedInFilter").asInt(), 0); + + response = postQuery("SELECT COUNT(*) FROM mytable WHERE REGEXP_LIKE(" + column + ", '.*')"); + assertEquals(response.get("resultTable").get("rows").get(0).get(0).asInt(), numTotalDocs); + assertEquals(response.get("numEntriesScannedInFilter").asInt(), 0); + } + } + private void testExpressionOverride() throws Exception { String query = "SELECT COUNT(*) FROM mytable WHERE DaysSinceEpoch * 24 = 392184"; From 221441f8165cb40c95ddb966149494620a686b2f Mon Sep 17 00:00:00 2001 From: vedant280 <87263239+vedant280@users.noreply.github.com> Date: Fri, 1 Aug 2025 16:08:54 +0530 Subject: [PATCH 059/167] Fix field shadowing in SparkSegmentTarPushJobRunner (#16482) Remove duplicate _spec field declaration that was shadowing the parent class field, causing null pointer issues. Fixes #16481 --- .../ingestion/batch/spark3/SparkSegmentTarPushJobRunner.java | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-3/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark3/SparkSegmentTarPushJobRunner.java b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-3/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark3/SparkSegmentTarPushJobRunner.java index a8c51fd606e3..4abfc449f398 100644 --- a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-3/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark3/SparkSegmentTarPushJobRunner.java +++ b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-3/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark3/SparkSegmentTarPushJobRunner.java @@ -36,8 +36,7 @@ import org.apache.spark.api.java.function.VoidFunction; public class SparkSegmentTarPushJobRunner extends BaseSparkSegmentTarPushJobRunner { - private SegmentGenerationJobSpec _spec; - + public SparkSegmentTarPushJobRunner() { super(); } From 1a4a5309c5d684df8a8a39051f388324531b76ab Mon Sep 17 00:00:00 2001 From: "Xiaotian (Jackie) Jiang" <17555551+Jackie-Jiang@users.noreply.github.com> Date: Fri, 1 Aug 2025 12:36:54 -0600 Subject: [PATCH 060/167] Cleanup deprecated methods in ThreadResourceUsageAccountant (#16479) --- .../restlet/resources/SystemResourceInfo.java | 7 +- .../pinot/common/utils/PinotAppConfigs.java | 13 ++-- .../HeapUsagePublishingAccountantFactory.java | 7 +- .../PerQueryCPUMemAccountantFactory.java | 42 ++---------- .../core/accounting/QueryAggregator.java | 11 ++- .../core/accounting/ResourceAggregator.java | 50 +++++++------- .../ResourceUsageAccountantFactory.java | 67 ++++--------------- .../core/accounting/WorkloadAggregator.java | 1 + .../query/scheduler/WorkloadScheduler.java | 2 +- .../ResourceManagerAccountingTest.java | 1 + .../core/accounting/TestThreadMXBean.java | 18 ++--- .../perf/BenchmarkWorkloadBudgetManager.java | 2 +- .../minion/tasks/purge/PurgeTaskExecutor.java | 11 +-- .../ThreadResourceUsageAccountant.java | 38 +---------- .../ThreadResourceUsageProvider.java | 13 ++-- .../accounting/WorkloadBudgetManager.java | 2 +- .../org/apache/pinot/spi/trace/Tracing.java | 62 ++--------------- .../pinot/spi/utils/ResourceUsageUtils.java | 43 ++++++++++++ .../accounting/WorkloadBudgetManagerTest.java | 6 +- ...rottleOnCriticalHeapUsageExecutorTest.java | 30 ++------- 20 files changed, 133 insertions(+), 293 deletions(-) rename pinot-spi/src/main/java/org/apache/pinot/{core => spi}/accounting/WorkloadBudgetManager.java (99%) create mode 100644 pinot-spi/src/main/java/org/apache/pinot/spi/utils/ResourceUsageUtils.java rename {pinot-core/src/test/java/org/apache/pinot/core => pinot-spi/src/test/java/org/apache/pinot/spi}/accounting/WorkloadBudgetManagerTest.java (97%) diff --git a/pinot-common/src/main/java/org/apache/pinot/common/restlet/resources/SystemResourceInfo.java b/pinot-common/src/main/java/org/apache/pinot/common/restlet/resources/SystemResourceInfo.java index b188ad3f1adb..8220a12ed460 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/restlet/resources/SystemResourceInfo.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/restlet/resources/SystemResourceInfo.java @@ -20,11 +20,10 @@ import com.google.common.collect.ImmutableMap; import java.lang.management.ManagementFactory; -import java.lang.management.MemoryMXBean; -import java.lang.management.MemoryUsage; import java.lang.management.OperatingSystemMXBean; import java.util.HashMap; import java.util.Map; +import org.apache.pinot.spi.utils.ResourceUsageUtils; /** @@ -58,9 +57,7 @@ public SystemResourceInfo() { _totalMemoryMB = runtime.totalMemory() / MEGA_BYTES; } - MemoryMXBean memoryMXBean = ManagementFactory.getMemoryMXBean(); - MemoryUsage heapMemoryUsage = memoryMXBean.getHeapMemoryUsage(); - _maxHeapSizeMB = heapMemoryUsage.getMax() / MEGA_BYTES; + _maxHeapSizeMB = ResourceUsageUtils.getMaxHeapSize() / MEGA_BYTES; } /** diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/PinotAppConfigs.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/PinotAppConfigs.java index 0f29c3a28711..67c36e719dcc 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/PinotAppConfigs.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/PinotAppConfigs.java @@ -25,20 +25,20 @@ import com.fasterxml.jackson.core.JsonProcessingException; import java.lang.management.GarbageCollectorMXBean; import java.lang.management.ManagementFactory; -import java.lang.management.MemoryMXBean; import java.lang.management.MemoryManagerMXBean; import java.lang.management.MemoryUsage; import java.lang.management.OperatingSystemMXBean; import java.lang.management.RuntimeMXBean; -import java.lang.management.ThreadMXBean; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.stream.Collectors; import org.apache.commons.io.FileUtils; +import org.apache.pinot.spi.accounting.ThreadResourceUsageProvider; import org.apache.pinot.spi.env.PinotConfiguration; import org.apache.pinot.spi.utils.JsonUtils; import org.apache.pinot.spi.utils.Obfuscator; +import org.apache.pinot.spi.utils.ResourceUsageUtils; /** @@ -126,12 +126,9 @@ private JVMConfig buildJVMConfig() { } private RuntimeConfig buildRuntimeConfig() { - MemoryMXBean memoryMXBean = ManagementFactory.getMemoryMXBean(); - MemoryUsage heapMemoryUsage = memoryMXBean.getHeapMemoryUsage(); - - ThreadMXBean threadMXBean = ManagementFactory.getThreadMXBean(); - return new RuntimeConfig(threadMXBean.getTotalStartedThreadCount(), threadMXBean.getThreadCount(), - FileUtils.byteCountToDisplaySize(heapMemoryUsage.getMax()), + MemoryUsage heapMemoryUsage = ResourceUsageUtils.getHeapMemoryUsage(); + return new RuntimeConfig(ThreadResourceUsageProvider.getTotalStartedThreadCount(), + ThreadResourceUsageProvider.getThreadCount(), FileUtils.byteCountToDisplaySize(heapMemoryUsage.getMax()), FileUtils.byteCountToDisplaySize(heapMemoryUsage.getUsed())); } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/accounting/HeapUsagePublishingAccountantFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/accounting/HeapUsagePublishingAccountantFactory.java index 7e82550286ea..c9861b74df39 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/accounting/HeapUsagePublishingAccountantFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/accounting/HeapUsagePublishingAccountantFactory.java @@ -18,8 +18,6 @@ */ package org.apache.pinot.core.accounting; -import java.lang.management.ManagementFactory; -import java.lang.management.MemoryMXBean; import java.util.Timer; import java.util.TimerTask; import org.apache.pinot.common.metrics.ServerGauge; @@ -30,6 +28,7 @@ import org.apache.pinot.spi.env.PinotConfiguration; import org.apache.pinot.spi.trace.Tracing; import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.ResourceUsageUtils; /** @@ -46,7 +45,6 @@ public ThreadResourceUsageAccountant init(PinotConfiguration config, String inst } public static class HeapUsagePublishingResourceUsageAccountant extends Tracing.DefaultThreadResourceUsageAccountant { - static final MemoryMXBean MEMORY_MX_BEAN = ManagementFactory.getMemoryMXBean(); private final Timer _timer; private final int _period; @@ -56,8 +54,7 @@ public HeapUsagePublishingResourceUsageAccountant(int period) { } public void publishHeapUsageMetrics() { - ServerMetrics.get() - .setValueOfGlobalGauge(ServerGauge.JVM_HEAP_USED_BYTES, MEMORY_MX_BEAN.getHeapMemoryUsage().getUsed()); + ServerMetrics.get().setValueOfGlobalGauge(ServerGauge.JVM_HEAP_USED_BYTES, ResourceUsageUtils.getUsedHeapSize()); } @Override diff --git a/pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java index 46e700fe8438..b88f06f2917d 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java @@ -21,8 +21,6 @@ import com.fasterxml.jackson.annotation.JsonIgnore; import com.google.common.cache.Cache; import com.google.common.cache.CacheBuilder; -import java.lang.management.ManagementFactory; -import java.lang.management.MemoryMXBean; import java.util.Collection; import java.util.Collections; import java.util.Comparator; @@ -57,6 +55,7 @@ import org.apache.pinot.spi.env.PinotConfiguration; import org.apache.pinot.spi.metrics.PinotMetricUtils; import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.ResourceUsageUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -78,11 +77,6 @@ public ThreadResourceUsageAccountant init(PinotConfiguration config, String inst } public static class PerQueryCPUMemResourceUsageAccountant implements ThreadResourceUsageAccountant { - - /** - * MemoryMXBean to get total heap used memory - */ - static final MemoryMXBean MEMORY_MX_BEAN = ManagementFactory.getMemoryMXBean(); private static final Logger LOGGER = LoggerFactory.getLogger(PerQueryCPUMemResourceUsageAccountant.class); private static final boolean IS_DEBUG_MODE_ENABLED = LOGGER.isDebugEnabled(); /** @@ -326,22 +320,6 @@ public boolean isQueryTerminated() { return false; } - @Override - @Deprecated - public void createExecutionContext(String queryId, int taskId, ThreadExecutionContext.TaskType taskType, - @Nullable ThreadExecutionContext parentContext) { - } - - @Override - @Deprecated - public void setThreadResourceUsageProvider(ThreadResourceUsageProvider threadResourceUsageProvider) { - } - - @Override - @Deprecated - public void updateQueryUsageConcurrently(String queryId, long cpuTimeNs, long memoryAllocatedBytes) { - } - @Override public void updateQueryUsageConcurrently(String queryId, long cpuTimeNs, long memoryAllocatedBytes, TrackingScope trackingScope) { @@ -359,11 +337,6 @@ public void updateQueryUsageConcurrently(String queryId, long cpuTimeNs, long me } } - @Override - @Deprecated - public void updateQueryUsageConcurrently(String queryId) { - } - /** * The thread would need to do {@code setThreadResourceUsageProvider} first upon it is scheduled. * This is to be called from a worker or a runner thread to update its corresponding cpu usage entry @@ -386,13 +359,6 @@ public void sampleThreadBytesAllocated() { } } - @Deprecated - @Override - public void setupRunner(String queryId, int taskId, ThreadExecutionContext.TaskType taskType) { - setupRunner(queryId, taskId, taskType, CommonConstants.Accounting.DEFAULT_WORKLOAD_NAME); - } - - @Override public void setupRunner(@Nullable String queryId, int taskId, ThreadExecutionContext.TaskType taskType, String workloadName) { @@ -559,8 +525,8 @@ public void reapFinishedTasks() { Thread thread = entry.getKey(); if (!thread.isAlive()) { + LOGGER.debug("Thread: {} is no longer alive, removing it from _threadEntriesMap", thread.getName()); _threadEntriesMap.remove(thread); - LOGGER.debug("Removing thread from _threadLocalEntry: {}", thread.getName()); } } _cancelSentQueries = cancellingQueries; @@ -701,7 +667,7 @@ public class WatcherTask implements Runnable, PinotClusterConfigChangeListener { private final AbstractMetrics.Gauge _memoryUsageGauge; WatcherTask() { - _queryMonitorConfig.set(new QueryMonitorConfig(_config, MEMORY_MX_BEAN.getHeapMemoryUsage().getMax())); + _queryMonitorConfig.set(new QueryMonitorConfig(_config, ResourceUsageUtils.getMaxHeapSize())); logQueryMonitorConfig(); switch (_instanceType) { @@ -835,7 +801,7 @@ public void runOnce() { } private void collectTriggerMetrics() { - _usedBytes = MEMORY_MX_BEAN.getHeapMemoryUsage().getUsed(); + _usedBytes = ResourceUsageUtils.getUsedHeapSize(); LOGGER.debug("Heap used bytes {}", _usedBytes); } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/accounting/QueryAggregator.java b/pinot-core/src/main/java/org/apache/pinot/core/accounting/QueryAggregator.java index 7c8bb815109b..4c051c1b5b8b 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/accounting/QueryAggregator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/accounting/QueryAggregator.java @@ -19,8 +19,6 @@ package org.apache.pinot.core.accounting; import com.fasterxml.jackson.annotation.JsonIgnore; -import java.lang.management.ManagementFactory; -import java.lang.management.MemoryMXBean; import java.util.Collections; import java.util.Comparator; import java.util.HashMap; @@ -42,6 +40,7 @@ import org.apache.pinot.spi.env.PinotConfiguration; import org.apache.pinot.spi.metrics.PinotMetricUtils; import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.ResourceUsageUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -56,8 +55,6 @@ public class QueryAggregator implements ResourceAggregator { private static final Logger LOGGER = LoggerFactory.getLogger(QueryAggregator.class); - static final MemoryMXBean MEMORY_MX_BEAN = ManagementFactory.getMemoryMXBean(); - enum TriggeringLevel { Normal, HeapMemoryAlarmingVerbose, CPUTimeBasedKilling, HeapMemoryCritical, HeapMemoryPanic } @@ -81,7 +78,7 @@ enum TriggeringLevel { private final String _instanceId; // max heap usage, Xmx - private final long _maxHeapSize = MEMORY_MX_BEAN.getHeapMemoryUsage().getMax(); + private final long _maxHeapSize = ResourceUsageUtils.getMaxHeapSize(); // don't kill a query if its memory footprint is below some ratio of _maxHeapSize private final long _minMemoryFootprintForKill; @@ -420,7 +417,7 @@ private void killMostExpensiveQuery() { Thread.sleep(_gcWaitTime); } catch (InterruptedException ignored) { } - _usedBytes = MEMORY_MX_BEAN.getHeapMemoryUsage().getUsed(); + _usedBytes = ResourceUsageUtils.getUsedHeapSize(); if (_usedBytes < _criticalLevelAfterGC) { return; } @@ -637,7 +634,7 @@ public void cleanUpPostAggregation() { } private void collectTriggerMetrics() { - _usedBytes = MEMORY_MX_BEAN.getHeapMemoryUsage().getUsed(); + _usedBytes = ResourceUsageUtils.getUsedHeapSize(); LOGGER.debug("Heap used bytes {}", _usedBytes); } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/accounting/ResourceAggregator.java b/pinot-core/src/main/java/org/apache/pinot/core/accounting/ResourceAggregator.java index 59649c1061af..30566e2f1f49 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/accounting/ResourceAggregator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/accounting/ResourceAggregator.java @@ -21,39 +21,35 @@ import java.util.List; -/** - * Interface for aggregating CPU and memory usage of threads. - */ +/// Interface for aggregating CPU and memory usage of threads. public interface ResourceAggregator { - /** - * Update CPU usage for one-off cases where identifier is known before-hand. For example: broker inbound netty - * thread where queryId and workloadName are already known. - * - * @param name identifier name - workload name, queryId, etc. - * @param cpuTimeNs CPU time in nanoseconds - */ - public void updateConcurrentCpuUsage(String name, long cpuTimeNs); + /// Updates CPU usage for one-off cases where identifier is known beforehand. For example: broker inbound netty thread + /// where queryId and workloadName are already known. + /// + /// @param name identifier name - queryId, workload name, etc. + /// @param cpuTimeNs CPU time in nanoseconds + void updateConcurrentCpuUsage(String name, long cpuTimeNs); - /** - * Update CPU usage for one-off cases where identifier is known before-hand. For example: broker inbound netty - * @param name identifier name - workload name, queryId, etc. - * @param memBytes memory usage in bytes - */ - public void updateConcurrentMemUsage(String name, long memBytes); + /// Updates memory usage for one-off cases where identifier is known beforehand. For example: broker inbound netty + /// thread where queryId and workloadName are already known. + /// + /// @param name identifier name - queryId, workload name, etc. + /// @param memBytes memory usage in bytes + void updateConcurrentMemUsage(String name, long memBytes); - // Cleanup of state after periodic aggregation is complete. - public void cleanUpPostAggregation(); + /// Cleans up state after periodic aggregation is complete. + void cleanUpPostAggregation(); - // Sleep time between aggregations. - public int getAggregationSleepTimeMs(); + /// Sleep time between aggregations. + int getAggregationSleepTimeMs(); - // Pre-aggregation step to be called before the aggregation of all thread entries. - public void preAggregate(List anchorThreadEntries); + /// Pre-aggregation step to be called before the aggregation of all thread entries. + void preAggregate(List anchorThreadEntries); - // Aggregation of each thread entry - public void aggregate(Thread thread, CPUMemThreadLevelAccountingObjects.ThreadEntry threadEntry); + /// Aggregates on a thread entry. + void aggregate(Thread thread, CPUMemThreadLevelAccountingObjects.ThreadEntry threadEntry); - // Post-aggregation step to be called after the aggregation of all thread entries. - public void postAggregate(); + /// Post-aggregation step to be called after the aggregation of all thread entries. + void postAggregate(); } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/accounting/ResourceUsageAccountantFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/accounting/ResourceUsageAccountantFactory.java index 3b3d509bf5d7..e00f74d888d6 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/accounting/ResourceUsageAccountantFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/accounting/ResourceUsageAccountantFactory.java @@ -20,8 +20,7 @@ import java.util.ArrayList; import java.util.Collection; -import java.util.Collections; -import java.util.HashMap; +import java.util.EnumMap; import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; @@ -42,6 +41,7 @@ import org.slf4j.LoggerFactory; +// TODO: Incorporate query OOM kill handling in PerQueryCPUMemAccountantFactory into this class public class ResourceUsageAccountantFactory implements ThreadAccountantFactory { @Override @@ -62,8 +62,6 @@ public static class ResourceUsageAccountant implements ThreadResourceUsageAccoun return thread; }); - private final PinotConfiguration _config; - // the map to track stats entry for each thread, the entry will automatically be added when one calls // setThreadResourceUsageProvider on the thread, including but not limited to // server worker thread, runner thread, broker jetty thread, or broker netty thread @@ -87,19 +85,12 @@ public static class ResourceUsageAccountant implements ThreadResourceUsageAccoun // is sampling allowed for MSE queries private final boolean _isThreadSamplingEnabledForMSE; - // instance id of the current instance, for logging purpose - private final String _instanceId; - private final WatcherTask _watcherTask; - private final Map _resourceAggregators; - - private final InstanceType _instanceType; + private final EnumMap _resourceAggregators; public ResourceUsageAccountant(PinotConfiguration config, String instanceId, InstanceType instanceType) { LOGGER.info("Initializing ResourceUsageAccountant"); - _config = config; - _instanceId = instanceId; boolean threadCpuTimeMeasurementEnabled = ThreadResourceUsageProvider.isThreadCpuTimeMeasurementEnabled(); boolean threadMemoryMeasurementEnabled = ThreadResourceUsageProvider.isThreadMemoryMeasurementEnabled(); @@ -113,7 +104,6 @@ public ResourceUsageAccountant(PinotConfiguration config, String instanceId, Ins CommonConstants.Accounting.DEFAULT_ENABLE_THREAD_MEMORY_SAMPLING); LOGGER.info("cpuSamplingConfig: {}, memorySamplingConfig: {}", cpuSamplingConfig, memorySamplingConfig); - _instanceType = instanceType; _isThreadCPUSamplingEnabled = cpuSamplingConfig && threadCpuTimeMeasurementEnabled; _isThreadMemorySamplingEnabled = memorySamplingConfig && threadMemoryMeasurementEnabled; LOGGER.info("_isThreadCPUSamplingEnabled: {}, _isThreadMemorySamplingEnabled: {}", _isThreadCPUSamplingEnabled, @@ -126,15 +116,15 @@ public ResourceUsageAccountant(PinotConfiguration config, String instanceId, Ins _watcherTask = new WatcherTask(); - _resourceAggregators = new HashMap<>(); + _resourceAggregators = new EnumMap<>(TrackingScope.class); // Add all aggregators. Configs of enabling/disabling cost collection/enforcement are handled in the aggregators. _resourceAggregators.put(TrackingScope.WORKLOAD, - new WorkloadAggregator(_isThreadCPUSamplingEnabled, _isThreadMemorySamplingEnabled, _config, _instanceType, - _instanceId)); + new WorkloadAggregator(_isThreadCPUSamplingEnabled, _isThreadMemorySamplingEnabled, config, instanceType, + instanceId)); _resourceAggregators.put(TrackingScope.QUERY, - new QueryAggregator(_isThreadCPUSamplingEnabled, _isThreadMemorySamplingEnabled, _config, _instanceType, - _instanceId)); + new QueryAggregator(_isThreadCPUSamplingEnabled, _isThreadMemorySamplingEnabled, config, instanceType, + instanceId)); } @Override @@ -166,32 +156,6 @@ public boolean isAnchorThreadInterrupted() { return false; } - @Override - @Deprecated - public void createExecutionContext(String queryId, int taskId, ThreadExecutionContext.TaskType taskType, - @Nullable ThreadExecutionContext parentContext) { - } - - @Override - @Deprecated - public void setThreadResourceUsageProvider(ThreadResourceUsageProvider threadResourceUsageProvider) { - } - - @Override - @Deprecated - public void updateQueryUsageConcurrently(String queryId, long cpuTimeNs, long memoryAllocatedBytes) { - } - - @Override - @Deprecated - public void updateQueryUsageConcurrently(String queryId) { - } - - @Override - public void setupRunner(String queryId, int taskId, ThreadExecutionContext.TaskType taskType) { - setupRunner(queryId, taskId, taskType, CommonConstants.Accounting.DEFAULT_WORKLOAD_NAME); - } - @Override public void setupRunner(String queryId, int taskId, ThreadExecutionContext.TaskType taskType, String workloadName) { _threadLocalEntry.get()._errorStatus.set(null); @@ -204,11 +168,12 @@ public void setupRunner(String queryId, int taskId, ThreadExecutionContext.TaskT @Override public void setupWorker(int taskId, ThreadExecutionContext.TaskType taskType, - @Nullable ThreadExecutionContext parentContext) { + @Nullable ThreadExecutionContext parentContext) { _threadLocalEntry.get()._errorStatus.set(null); if (parentContext != null && parentContext.getQueryId() != null && parentContext.getAnchorThread() != null) { - _threadLocalEntry.get().setThreadTaskStatus(parentContext.getQueryId(), taskId, parentContext.getTaskType(), - parentContext.getAnchorThread(), parentContext.getWorkloadName()); + _threadLocalEntry.get() + .setThreadTaskStatus(parentContext.getQueryId(), taskId, parentContext.getTaskType(), + parentContext.getAnchorThread(), parentContext.getWorkloadName()); } } @@ -226,10 +191,6 @@ public int getEntryCount() { @Override public Map getQueryResources() { QueryAggregator queryAggregator = (QueryAggregator) _resourceAggregators.get(TrackingScope.QUERY); - if (queryAggregator == null) { - return Collections.emptyMap(); - } - return queryAggregator.getQueryResources(_threadEntriesMap); } @@ -237,9 +198,6 @@ public int getEntryCount() { public void updateQueryUsageConcurrently(String identifier, long cpuTimeNs, long memoryAllocatedBytes, TrackingScope trackingScope) { ResourceAggregator resourceAggregator = _resourceAggregators.get(trackingScope); - if (resourceAggregator == null) { - return; - } if (_isThreadCPUSamplingEnabled) { resourceAggregator.updateConcurrentCpuUsage(identifier, cpuTimeNs); } @@ -312,7 +270,6 @@ public List getAnchorThreadEntri return anchorThreadEntries; } - class WatcherTask implements Runnable { WatcherTask() { } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/accounting/WorkloadAggregator.java b/pinot-core/src/main/java/org/apache/pinot/core/accounting/WorkloadAggregator.java index 33de4bf2a1b7..72d0d818cda4 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/accounting/WorkloadAggregator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/accounting/WorkloadAggregator.java @@ -23,6 +23,7 @@ import java.util.List; import java.util.Map; import java.util.concurrent.ConcurrentHashMap; +import org.apache.pinot.spi.accounting.WorkloadBudgetManager; import org.apache.pinot.spi.config.instance.InstanceType; import org.apache.pinot.spi.env.PinotConfiguration; import org.apache.pinot.spi.trace.Tracing; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/WorkloadScheduler.java b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/WorkloadScheduler.java index 9d0fbb062154..e90f624e22e9 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/WorkloadScheduler.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/scheduler/WorkloadScheduler.java @@ -26,12 +26,12 @@ import org.apache.pinot.common.metrics.ServerMetrics; import org.apache.pinot.common.metrics.ServerQueryPhase; import org.apache.pinot.common.utils.config.QueryOptionsUtils; -import org.apache.pinot.core.accounting.WorkloadBudgetManager; import org.apache.pinot.core.query.executor.QueryExecutor; import org.apache.pinot.core.query.request.ServerQueryRequest; import org.apache.pinot.core.query.scheduler.resources.QueryExecutorService; import org.apache.pinot.core.query.scheduler.resources.UnboundedResourceManager; import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant; +import org.apache.pinot.spi.accounting.WorkloadBudgetManager; import org.apache.pinot.spi.env.PinotConfiguration; import org.apache.pinot.spi.query.QueryThreadContext; import org.apache.pinot.spi.trace.Tracing; diff --git a/pinot-core/src/test/java/org/apache/pinot/core/accounting/ResourceManagerAccountingTest.java b/pinot-core/src/test/java/org/apache/pinot/core/accounting/ResourceManagerAccountingTest.java index 4986e7684d35..2641cd60bd3b 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/accounting/ResourceManagerAccountingTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/accounting/ResourceManagerAccountingTest.java @@ -61,6 +61,7 @@ import org.apache.pinot.spi.accounting.ThreadExecutionContext; import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant; import org.apache.pinot.spi.accounting.ThreadResourceUsageProvider; +import org.apache.pinot.spi.accounting.WorkloadBudgetManager; import org.apache.pinot.spi.config.instance.InstanceType; import org.apache.pinot.spi.config.table.JsonIndexConfig; import org.apache.pinot.spi.env.PinotConfiguration; diff --git a/pinot-core/src/test/java/org/apache/pinot/core/accounting/TestThreadMXBean.java b/pinot-core/src/test/java/org/apache/pinot/core/accounting/TestThreadMXBean.java index 572c2afaa8ed..6d5614f43e04 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/accounting/TestThreadMXBean.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/accounting/TestThreadMXBean.java @@ -18,8 +18,6 @@ */ package org.apache.pinot.core.accounting; -import java.lang.management.ManagementFactory; -import java.lang.management.MemoryMXBean; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; @@ -28,6 +26,7 @@ import org.apache.log4j.LogManager; import org.apache.pinot.spi.accounting.ThreadResourceSnapshot; import org.apache.pinot.spi.accounting.ThreadResourceUsageProvider; +import org.apache.pinot.spi.utils.ResourceUsageUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.testng.Assert; @@ -72,10 +71,9 @@ public void testThreadMXBeanMultithreadMemAllocTracking() { AtomicLong b = new AtomicLong(); AtomicLong c = new AtomicLong(); ExecutorService executor = Executors.newFixedThreadPool(3); - MemoryMXBean memoryMXBean = ManagementFactory.getMemoryMXBean(); System.gc(); - long heapPrev = memoryMXBean.getHeapMemoryUsage().getUsed(); + long heapPrev = ResourceUsageUtils.getUsedHeapSize(); ThreadResourceSnapshot threadResourceSnapshot0 = new ThreadResourceSnapshot(); executor.submit(() -> { ThreadResourceSnapshot threadResourceSnapshot = new ThreadResourceSnapshot(); @@ -108,7 +106,7 @@ public void testThreadMXBeanMultithreadMemAllocTracking() { long d = threadResourceSnapshot0.getAllocatedBytes(); long threadAllocatedBytes = a.get() + b.get() + c.get() + d; - float heapUsedBytes = (float) memoryMXBean.getHeapMemoryUsage().getUsed() - heapPrev; + float heapUsedBytes = (float) ResourceUsageUtils.getUsedHeapSize() - heapPrev; float ratio = threadAllocatedBytes / heapUsedBytes; LOGGER.info("Measured thread allocated bytes {}, heap used bytes {}, ratio {}", @@ -129,10 +127,9 @@ public void testThreadMXBeanDeepMemAllocTracking() { AtomicLong b = new AtomicLong(); AtomicLong c = new AtomicLong(); ExecutorService executor = Executors.newFixedThreadPool(3); - MemoryMXBean memoryMXBean = ManagementFactory.getMemoryMXBean(); System.gc(); - long heapPrev = memoryMXBean.getHeapMemoryUsage().getUsed(); + long heapPrev = ResourceUsageUtils.getUsedHeapSize(); ThreadResourceSnapshot threadResourceSnapshot0 = new ThreadResourceSnapshot(); executor.submit(() -> { ThreadResourceSnapshot threadResourceSnapshot = new ThreadResourceSnapshot(); @@ -165,7 +162,7 @@ public void testThreadMXBeanDeepMemAllocTracking() { long d = threadResourceSnapshot0.getAllocatedBytes(); long threadAllocatedBytes = a.get() + b.get() + c.get() + d; - float heapUsedBytes = (float) memoryMXBean.getHeapMemoryUsage().getUsed() - heapPrev; + float heapUsedBytes = (float) ResourceUsageUtils.getUsedHeapSize() - heapPrev; float ratio = threadAllocatedBytes / heapUsedBytes; LOGGER.info("Measured thread allocated bytes {}, heap used bytes {}, ratio {}", @@ -180,15 +177,14 @@ public void testThreadMXBeanDeepMemAllocTracking() { @SuppressWarnings("unused") public void testThreadMXBeanMemAllocGCTracking() { LogManager.getLogger(TestThreadMXBean.class).setLevel(Level.INFO); - MemoryMXBean memoryMXBean = ManagementFactory.getMemoryMXBean(); System.gc(); ThreadResourceSnapshot threadResourceSnapshot0 = new ThreadResourceSnapshot(); - long heapPrev = memoryMXBean.getHeapMemoryUsage().getUsed(); + long heapPrev = ResourceUsageUtils.getUsedHeapSize(); for (int i = 0; i < 3; i++) { long[] ignored = new long[100000000]; } System.gc(); - long heapResult = memoryMXBean.getHeapMemoryUsage().getUsed() - heapPrev; + long heapResult = ResourceUsageUtils.getUsedHeapSize() - heapPrev; long result = threadResourceSnapshot0.getAllocatedBytes(); LOGGER.info("Measured thread allocated bytes {}, heap used bytes {}", result, heapResult); diff --git a/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkWorkloadBudgetManager.java b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkWorkloadBudgetManager.java index 05805c2bb32a..3904920ec283 100644 --- a/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkWorkloadBudgetManager.java +++ b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkWorkloadBudgetManager.java @@ -19,7 +19,7 @@ package org.apache.pinot.perf; import java.util.concurrent.TimeUnit; -import org.apache.pinot.core.accounting.WorkloadBudgetManager; +import org.apache.pinot.spi.accounting.WorkloadBudgetManager; import org.apache.pinot.spi.env.PinotConfiguration; import org.apache.pinot.spi.utils.CommonConstants; import org.openjdk.jmh.annotations.Benchmark; diff --git a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/purge/PurgeTaskExecutor.java b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/purge/PurgeTaskExecutor.java index 58b0d16c2b4e..2877f1b3ccd2 100644 --- a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/purge/PurgeTaskExecutor.java +++ b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/purge/PurgeTaskExecutor.java @@ -19,8 +19,6 @@ package org.apache.pinot.plugin.minion.tasks.purge; import java.io.File; -import java.lang.management.ManagementFactory; -import java.lang.management.ThreadMXBean; import java.util.Collections; import java.util.Map; import java.util.concurrent.TimeUnit; @@ -32,21 +30,18 @@ import org.apache.pinot.core.minion.SegmentPurger; import org.apache.pinot.plugin.minion.tasks.BaseSingleSegmentConversionExecutor; import org.apache.pinot.plugin.minion.tasks.SegmentConversionResult; +import org.apache.pinot.spi.accounting.ThreadResourceUsageProvider; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.data.Schema; import org.apache.pinot.spi.utils.builder.TableNameBuilder; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; public class PurgeTaskExecutor extends BaseSingleSegmentConversionExecutor { - private static final Logger LOGGER = LoggerFactory.getLogger(PurgeTaskExecutor.class); protected final MinionMetrics _minionMetrics = MinionMetrics.get(); public static final String RECORD_PURGER_KEY = "recordPurger"; public static final String RECORD_MODIFIER_KEY = "recordModifier"; public static final String NUM_RECORDS_PURGED_KEY = "numRecordsPurged"; public static final String NUM_RECORDS_MODIFIED_KEY = "numRecordsModified"; - private static final ThreadMXBean MX_BEAN = ManagementFactory.getThreadMXBean(); @Override protected SegmentConversionResult convert(PinotTaskConfig pinotTaskConfig, File indexDir, File workingDir) @@ -68,9 +63,9 @@ protected SegmentConversionResult convert(PinotTaskConfig pinotTaskConfig, File _eventObserver.notifyProgress(pinotTaskConfig, "Purging segment: " + indexDir); SegmentPurger segmentPurger = new SegmentPurger(indexDir, workingDir, tableConfig, schema, recordPurger, recordModifier, null); - long purgeTaskStartTimeNs = MX_BEAN.getCurrentThreadCpuTime(); + long purgeTaskStartTimeNs = ThreadResourceUsageProvider.getCurrentThreadCpuTime(); File purgedSegmentFile = segmentPurger.purgeSegment(); - long purgeTaskEndTimeNs = MX_BEAN.getCurrentThreadCpuTime(); + long purgeTaskEndTimeNs = ThreadResourceUsageProvider.getCurrentThreadCpuTime(); _minionMetrics.addTimedTableValue(tableNameWithType, taskType, MinionTimer.TASK_THREAD_CPU_TIME_NS, purgeTaskEndTimeNs - purgeTaskStartTimeNs, TimeUnit.NANOSECONDS); if (purgedSegmentFile == null) { diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceUsageAccountant.java b/pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceUsageAccountant.java index dc912371980f..db41723d5bea 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceUsageAccountant.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceUsageAccountant.java @@ -46,23 +46,6 @@ default boolean isQueryTerminated() { return false; } - /** - * This method has been deprecated and replaced by {@link setupRunner(String, int, ThreadExecutionContext.TaskType)} - * and {@link setupWorker(int, ThreadExecutionContext.TaskType, ThreadExecutionContext)}. - */ - @Deprecated - void createExecutionContext(String queryId, int taskId, ThreadExecutionContext.TaskType taskType, - @Nullable ThreadExecutionContext parentContext); - - /** - * Set up the thread execution context for an anchor a.k.a runner thread. - * @param queryId query id string - * @param taskId a unique task id - * @param taskType the type of the task - SSE or MSE - */ - @Deprecated - void setupRunner(String queryId, int taskId, ThreadExecutionContext.TaskType taskType); - /** * Set up the thread execution context for an anchor a.k.a runner thread. * @param queryId query id string @@ -79,7 +62,7 @@ void createExecutionContext(String queryId, int taskId, ThreadExecutionContext.T * @param parentContext the parent execution context */ void setupWorker(int taskId, ThreadExecutionContext.TaskType taskType, - @Nullable ThreadExecutionContext parentContext); + @Nullable ThreadExecutionContext parentContext); /** * get the executon context of current thread @@ -87,12 +70,6 @@ void setupWorker(int taskId, ThreadExecutionContext.TaskType taskType, @Nullable ThreadExecutionContext getThreadExecutionContext(); - /** - * set resource usage provider - */ - @Deprecated - void setThreadResourceUsageProvider(ThreadResourceUsageProvider threadResourceUsageProvider); - /** * call to sample usage */ @@ -117,24 +94,13 @@ default void registerMseCancelCallback(String queryId, MseCancelCallback mseCanc // Default implementation does nothing. Subclasses can override to register a cancel callback. } - /** - * special interface to aggregate usage to the stats store only once, it is used for response - * ser/de threads where the thread execution context cannot be setup before hands as - * queryId/taskId is unknown and the execution process is hard to instrument - */ - @Deprecated - void updateQueryUsageConcurrently(String queryId, long cpuTimeNs, long allocatedBytes); - /** * special interface to aggregate usage to the stats store only once, it is used for response * ser/de threads where the thread execution context cannot be setup before hands as * queryId/taskId/workloadName is unknown and the execution process is hard to instrument */ void updateQueryUsageConcurrently(String identifier, long cpuTimeNs, long allocatedBytes, - TrackingScope trackingScope); - - @Deprecated - void updateQueryUsageConcurrently(String queryId); + TrackingScope trackingScope); /** * start the periodical task diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceUsageProvider.java b/pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceUsageProvider.java index f5ebc5ed2fa7..44fb9f6cea35 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceUsageProvider.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceUsageProvider.java @@ -31,6 +31,9 @@ * and allocateBytes (JVM heap) for the current thread. */ public class ThreadResourceUsageProvider { + private ThreadResourceUsageProvider() { + } + private static final Logger LOGGER = LoggerFactory.getLogger(ThreadResourceUsageProvider.class); // used for getting the memory allocation function in hotspot jvm through reflection @@ -51,14 +54,12 @@ public class ThreadResourceUsageProvider { private static boolean _isThreadCpuTimeMeasurementEnabled = false; private static boolean _isThreadMemoryMeasurementEnabled = false; - @Deprecated - public long getThreadTimeNs() { - return 0; + public static int getThreadCount() { + return MX_BEAN.getThreadCount(); } - @Deprecated - public long getThreadAllocatedBytes() { - return 0; + public static long getTotalStartedThreadCount() { + return MX_BEAN.getTotalStartedThreadCount(); } public static long getCurrentThreadCpuTime() { diff --git a/pinot-spi/src/main/java/org/apache/pinot/core/accounting/WorkloadBudgetManager.java b/pinot-spi/src/main/java/org/apache/pinot/spi/accounting/WorkloadBudgetManager.java similarity index 99% rename from pinot-spi/src/main/java/org/apache/pinot/core/accounting/WorkloadBudgetManager.java rename to pinot-spi/src/main/java/org/apache/pinot/spi/accounting/WorkloadBudgetManager.java index 5d0004dd628b..dc1642a374c3 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/core/accounting/WorkloadBudgetManager.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/accounting/WorkloadBudgetManager.java @@ -16,7 +16,7 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.pinot.core.accounting; +package org.apache.pinot.spi.accounting; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.Executors; diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/trace/Tracing.java b/pinot-spi/src/main/java/org/apache/pinot/spi/trace/Tracing.java index 0f2ee9deb725..6d9993cd4b71 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/trace/Tracing.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/trace/Tracing.java @@ -25,14 +25,13 @@ import java.util.Map; import java.util.concurrent.atomic.AtomicReference; import javax.annotation.Nullable; -import org.apache.pinot.core.accounting.WorkloadBudgetManager; import org.apache.pinot.spi.accounting.QueryResourceTracker; import org.apache.pinot.spi.accounting.ThreadAccountantFactory; import org.apache.pinot.spi.accounting.ThreadExecutionContext; import org.apache.pinot.spi.accounting.ThreadResourceTracker; import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant; -import org.apache.pinot.spi.accounting.ThreadResourceUsageProvider; import org.apache.pinot.spi.accounting.TrackingScope; +import org.apache.pinot.spi.accounting.WorkloadBudgetManager; import org.apache.pinot.spi.config.instance.InstanceType; import org.apache.pinot.spi.env.PinotConfiguration; import org.apache.pinot.spi.exception.EarlyTerminationException; @@ -199,17 +198,6 @@ public boolean isAnchorThreadInterrupted() { return false; } - @Override - public void createExecutionContext(String queryId, int taskId, ThreadExecutionContext.TaskType taskType, - @Nullable ThreadExecutionContext parentContext) { - } - - @Deprecated - public void createExecutionContextInner(@Nullable String queryId, int taskId, - ThreadExecutionContext.TaskType taskType, - @Nullable ThreadExecutionContext parentContext) { - } - @Override public void clear() { } @@ -222,38 +210,19 @@ public void sampleUsage() { public void sampleUsageMSE() { } - @Deprecated - public void setThreadResourceUsageProvider(ThreadResourceUsageProvider threadResourceUsageProvider) { - } - - @Override - @Deprecated - public void updateQueryUsageConcurrently(String queryId) { - // No-op for default accountant - } - - @Override - public void updateQueryUsageConcurrently(String queryId, long cpuTimeNs, long allocatedBytes) { - // No-op for default accountant - } - @Override public void updateQueryUsageConcurrently(String queryId, long cpuTimeNs, long allocatedBytes, - TrackingScope trackingScope) { + TrackingScope trackingScope) { // No-op for default accountant } - @Override - public void setupRunner(String queryId, int taskId, ThreadExecutionContext.TaskType taskType) { - } - @Override public void setupRunner(String queryId, int taskId, ThreadExecutionContext.TaskType taskType, String workloadName) { } @Override public void setupWorker(int taskId, ThreadExecutionContext.TaskType taskType, - @Nullable ThreadExecutionContext parentContext) { + @Nullable ThreadExecutionContext parentContext) { } @Override @@ -293,22 +262,14 @@ public static class ThreadAccountantOps { private ThreadAccountantOps() { } - @Deprecated - public static void setupRunner(String queryId) { - } - - @Deprecated - public static void setupRunner(String queryId, ThreadExecutionContext.TaskType taskType) { - } - public static void setupRunner(String queryId, String workloadName) { setupRunner(queryId, ThreadExecutionContext.TaskType.SSE, workloadName); } public static void setupRunner(String queryId, ThreadExecutionContext.TaskType taskType, String workloadName) { // Set up the runner thread with the given query ID and workload name - Tracing.getThreadAccountant().setupRunner(queryId, CommonConstants.Accounting.ANCHOR_TASK_ID, taskType, - workloadName); + Tracing.getThreadAccountant() + .setupRunner(queryId, CommonConstants.Accounting.ANCHOR_TASK_ID, taskType, workloadName); } /** @@ -398,24 +359,11 @@ public static void sampleAndCheckInterruption(ThreadResourceUsageAccountant acco accountant.sampleUsage(); } - @Deprecated - public static void updateQueryUsageConcurrently(String queryId) { - } - - @Deprecated - public static void updateQueryUsageConcurrently(String queryId, long cpuTimeNs, long allocatedBytes) { - Tracing.getThreadAccountant().updateQueryUsageConcurrently(queryId, cpuTimeNs, allocatedBytes); - } - public static void updateQueryUsageConcurrently(String queryId, long cpuTimeNs, long allocatedBytes, TrackingScope trackingScope) { Tracing.getThreadAccountant().updateQueryUsageConcurrently(queryId, cpuTimeNs, allocatedBytes, trackingScope); } - @Deprecated - public static void setThreadResourceUsageProvider() { - } - // Check for thread interruption, every time after merging 8192 keys public static void sampleAndCheckInterruptionPeriodically(int mergedKeys) { if ((mergedKeys & MAX_ENTRIES_KEYS_MERGED_PER_INTERRUPTION_CHECK_MASK) == 0) { diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/ResourceUsageUtils.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/ResourceUsageUtils.java new file mode 100644 index 000000000000..d6ba26b0e4ee --- /dev/null +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/ResourceUsageUtils.java @@ -0,0 +1,43 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.spi.utils; + +import java.lang.management.ManagementFactory; +import java.lang.management.MemoryMXBean; +import java.lang.management.MemoryUsage; + + +public class ResourceUsageUtils { + private ResourceUsageUtils() { + } + + private static final MemoryMXBean MEMORY_MX_BEAN = ManagementFactory.getMemoryMXBean(); + + public static MemoryUsage getHeapMemoryUsage() { + return MEMORY_MX_BEAN.getHeapMemoryUsage(); + } + + public static long getMaxHeapSize() { + return MEMORY_MX_BEAN.getHeapMemoryUsage().getMax(); + } + + public static long getUsedHeapSize() { + return MEMORY_MX_BEAN.getHeapMemoryUsage().getUsed(); + } +} diff --git a/pinot-core/src/test/java/org/apache/pinot/core/accounting/WorkloadBudgetManagerTest.java b/pinot-spi/src/test/java/org/apache/pinot/spi/accounting/WorkloadBudgetManagerTest.java similarity index 97% rename from pinot-core/src/test/java/org/apache/pinot/core/accounting/WorkloadBudgetManagerTest.java rename to pinot-spi/src/test/java/org/apache/pinot/spi/accounting/WorkloadBudgetManagerTest.java index 21cc2df4f0bc..d73ca4173a5e 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/accounting/WorkloadBudgetManagerTest.java +++ b/pinot-spi/src/test/java/org/apache/pinot/spi/accounting/WorkloadBudgetManagerTest.java @@ -16,7 +16,7 @@ * specific language governing permissions and limitations * under the License. */ -package org.apache.pinot.core.accounting; +package org.apache.pinot.spi.accounting; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; @@ -26,7 +26,9 @@ import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; -import static org.testng.Assert.*; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; public class WorkloadBudgetManagerTest { diff --git a/pinot-spi/src/test/java/org/apache/pinot/spi/executor/ThrottleOnCriticalHeapUsageExecutorTest.java b/pinot-spi/src/test/java/org/apache/pinot/spi/executor/ThrottleOnCriticalHeapUsageExecutorTest.java index bab8167e3f98..141b55129f0b 100644 --- a/pinot-spi/src/test/java/org/apache/pinot/spi/executor/ThrottleOnCriticalHeapUsageExecutorTest.java +++ b/pinot-spi/src/test/java/org/apache/pinot/spi/executor/ThrottleOnCriticalHeapUsageExecutorTest.java @@ -30,7 +30,6 @@ import org.apache.pinot.spi.accounting.ThreadExecutionContext; import org.apache.pinot.spi.accounting.ThreadResourceTracker; import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant; -import org.apache.pinot.spi.accounting.ThreadResourceUsageProvider; import org.apache.pinot.spi.accounting.TrackingScope; import org.apache.pinot.spi.exception.QueryErrorCode; import org.apache.pinot.spi.exception.QueryException; @@ -42,9 +41,11 @@ public class ThrottleOnCriticalHeapUsageExecutorTest { @Test - void testThrottle() throws Exception { + void testThrottle() + throws Exception { ThreadResourceUsageAccountant accountant = new ThreadResourceUsageAccountant() { final AtomicLong _numCalls = new AtomicLong(0); + @Override public void clear() { } @@ -54,15 +55,6 @@ public boolean isAnchorThreadInterrupted() { return false; } - @Override - public void createExecutionContext(String queryId, int taskId, ThreadExecutionContext.TaskType taskType, - @Nullable ThreadExecutionContext parentContext) { - } - - @Override - public void setupRunner(String queryId, int taskId, ThreadExecutionContext.TaskType taskType) { - } - @Override public void setupRunner(String queryId, int taskId, ThreadExecutionContext.TaskType taskType, String workloadName) { @@ -79,10 +71,6 @@ public ThreadExecutionContext getThreadExecutionContext() { return null; } - @Override - public void setThreadResourceUsageProvider(ThreadResourceUsageProvider threadResourceUsageProvider) { - } - @Override public void sampleUsage() { } @@ -96,14 +84,6 @@ public boolean throttleQuerySubmission() { return _numCalls.getAndIncrement() > 1; } - @Override - public void updateQueryUsageConcurrently(String queryId, long cpuTimeNs, long allocatedBytes) { - } - - @Override - public void updateQueryUsageConcurrently(String queryId) { - } - @Override public void updateQueryUsageConcurrently(String queryId, long cpuTimeNs, long allocatedBytes, TrackingScope trackingScope) { @@ -129,8 +109,8 @@ public Collection getThreadResources() { } }; - ThrottleOnCriticalHeapUsageExecutor executor = new ThrottleOnCriticalHeapUsageExecutor( - Executors.newCachedThreadPool(), accountant); + ThrottleOnCriticalHeapUsageExecutor executor = + new ThrottleOnCriticalHeapUsageExecutor(Executors.newCachedThreadPool(), accountant); CyclicBarrier barrier = new CyclicBarrier(2); From e59aa47c3525cdde3e0d8db3bb31fa874ab30cfb Mon Sep 17 00:00:00 2001 From: "Xiaotian (Jackie) Jiang" <17555551+Jackie-Jiang@users.noreply.github.com> Date: Fri, 1 Aug 2025 13:30:53 -0600 Subject: [PATCH 061/167] Fix linter violation in SparkSegmentTarPushJobRunner and also apply fix for spark2 (#16488) --- .../batch/spark/SparkSegmentTarPushJobRunner.java | 10 ++++------ .../batch/spark3/SparkSegmentTarPushJobRunner.java | 11 ++++++----- 2 files changed, 10 insertions(+), 11 deletions(-) diff --git a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-2.4/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark/SparkSegmentTarPushJobRunner.java b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-2.4/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark/SparkSegmentTarPushJobRunner.java index babe30e7690e..499ec2daba60 100644 --- a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-2.4/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark/SparkSegmentTarPushJobRunner.java +++ b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-2.4/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark/SparkSegmentTarPushJobRunner.java @@ -36,9 +36,7 @@ import org.apache.spark.api.java.function.VoidFunction; - public class SparkSegmentTarPushJobRunner extends BaseSparkSegmentTarPushJobRunner { - private SegmentGenerationJobSpec _spec; public SparkSegmentTarPushJobRunner() { super(); @@ -48,8 +46,8 @@ public SparkSegmentTarPushJobRunner(SegmentGenerationJobSpec spec) { super(spec); } - public void parallelizeTarPushJob(List pinotFSSpecs, - List segmentUris, int pushParallelism, URI outputDirURI) { + public void parallelizeTarPushJob(List pinotFSSpecs, List segmentUris, int pushParallelism, + URI outputDirURI) { JavaSparkContext sparkContext = JavaSparkContext.fromSparkContext(SparkContext.getOrCreate()); JavaRDD pathRDD = sparkContext.parallelize(segmentUris, pushParallelism); URI finalOutputDirURI = outputDirURI; @@ -61,8 +59,8 @@ public void call(String segmentTarPath) throws Exception { PluginManager.get().init(); for (PinotFSSpec pinotFSSpec : pinotFSSpecs) { - PinotFSFactory - .register(pinotFSSpec.getScheme(), pinotFSSpec.getClassName(), new PinotConfiguration(pinotFSSpec)); + PinotFSFactory.register(pinotFSSpec.getScheme(), pinotFSSpec.getClassName(), + new PinotConfiguration(pinotFSSpec)); } try { SegmentPushUtils.pushSegments(_spec, PinotFSFactory.create(finalOutputDirURI.getScheme()), diff --git a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-3/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark3/SparkSegmentTarPushJobRunner.java b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-3/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark3/SparkSegmentTarPushJobRunner.java index 4abfc449f398..ca27dd2b9835 100644 --- a/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-3/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark3/SparkSegmentTarPushJobRunner.java +++ b/pinot-plugins/pinot-batch-ingestion/pinot-batch-ingestion-spark-3/src/main/java/org/apache/pinot/plugin/ingestion/batch/spark3/SparkSegmentTarPushJobRunner.java @@ -35,8 +35,9 @@ import org.apache.spark.api.java.JavaSparkContext; import org.apache.spark.api.java.function.VoidFunction; + public class SparkSegmentTarPushJobRunner extends BaseSparkSegmentTarPushJobRunner { - + public SparkSegmentTarPushJobRunner() { super(); } @@ -45,8 +46,8 @@ public SparkSegmentTarPushJobRunner(SegmentGenerationJobSpec spec) { super(spec); } - public void parallelizeTarPushJob(List pinotFSSpecs, - List segmentUris, int pushParallelism, URI outputDirURI) { + public void parallelizeTarPushJob(List pinotFSSpecs, List segmentUris, int pushParallelism, + URI outputDirURI) { JavaSparkContext sparkContext = JavaSparkContext.fromSparkContext(SparkContext.getOrCreate()); JavaRDD pathRDD = sparkContext.parallelize(segmentUris, pushParallelism); URI finalOutputDirURI = outputDirURI; @@ -58,8 +59,8 @@ public void call(String segmentTarPath) throws Exception { PluginManager.get().init(); for (PinotFSSpec pinotFSSpec : pinotFSSpecs) { - PinotFSFactory - .register(pinotFSSpec.getScheme(), pinotFSSpec.getClassName(), new PinotConfiguration(pinotFSSpec)); + PinotFSFactory.register(pinotFSSpec.getScheme(), pinotFSSpec.getClassName(), + new PinotConfiguration(pinotFSSpec)); } try { SegmentPushUtils.pushSegments(_spec, PinotFSFactory.create(finalOutputDirURI.getScheme()), From 4cd3296060bb3cbc2131a8c72d89934506ed9b26 Mon Sep 17 00:00:00 2001 From: Shaurya Chaturvedi Date: Sat, 2 Aug 2025 09:36:44 -0700 Subject: [PATCH 062/167] [timeseries][bugfix] Fixed timeseries API clashes in Controller (#16487) Co-authored-by: Shaurya Chaturvedi --- .../PinotControllerTimeseriesResource.java | 4 +-- .../api/resources/PinotQueryResource.java | 2 +- .../pinot/integration/tests/ClusterTest.java | 11 ++++++- .../tests/TimeSeriesAuthIntegrationTest.java | 1 + .../tests/TimeSeriesIntegrationTest.java | 29 +++++++++++++++++++ 5 files changed, 43 insertions(+), 4 deletions(-) diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotControllerTimeseriesResource.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotControllerTimeseriesResource.java index 74c159c17336..2e2643f4ed25 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotControllerTimeseriesResource.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotControllerTimeseriesResource.java @@ -33,7 +33,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -@Path("/timeseries") +@Path("/") public class PinotControllerTimeseriesResource { public static final Logger LOGGER = LoggerFactory.getLogger(PinotControllerTimeseriesResource.class); @@ -42,7 +42,7 @@ public class PinotControllerTimeseriesResource { @GET @Produces(MediaType.APPLICATION_JSON) - @Path("languages") + @Path("/timeseries/languages") @ApiOperation(value = "Get timeseries languages from controller configuration", notes = "Get timeseries languages from controller configuration") public List getBrokerTimeSeriesLanguages(@Context HttpHeaders headers) { diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotQueryResource.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotQueryResource.java index 6676016c13ab..1e4d4f18863e 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotQueryResource.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotQueryResource.java @@ -136,7 +136,7 @@ public StreamingOutput handleGetSql(@QueryParam("sql") String sqlQuery, @QueryPa } @GET - @Path("timeseries/api/v1/query_range") + @Path("/timeseries/api/v1/query_range") @ManualAuthorization @ApiOperation(value = "Prometheus Compatible API for Pinot's Time Series Engine") public StreamingOutput handleTimeSeriesQueryRange(@QueryParam("language") String language, diff --git a/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/ClusterTest.java b/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/ClusterTest.java index a54011fc32cb..517f3bb8e547 100644 --- a/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/ClusterTest.java +++ b/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/ClusterTest.java @@ -560,10 +560,19 @@ public JsonNode postQuery(@Language("sql") String query) * This is used for testing timeseries queries. */ public JsonNode getTimeseriesQuery(String query, long startTime, long endTime, Map headers) { + return getTimeseriesQuery(getBrokerBaseApiUrl(), query, startTime, endTime, headers); + } + + /** + * Queries the timeseries query endpoint (/timeseries/api/v1/query_range) of the given base URL. + * This is used for testing timeseries queries. + */ + public JsonNode getTimeseriesQuery(String baseUrl, String query, long startTime, long endTime, + Map headers) { try { Map queryParams = Map.of("language", "m3ql", "query", query, "start", String.valueOf(startTime), "end", String.valueOf(endTime)); - String url = buildQueryUrl(getTimeSeriesQueryApiUrl(getBrokerBaseApiUrl()), queryParams); + String url = buildQueryUrl(getTimeSeriesQueryApiUrl(baseUrl), queryParams); JsonNode responseJsonNode = JsonUtils.stringToJsonNode(sendGetRequest(url, headers)); return sanitizeResponse(responseJsonNode); } catch (Exception e) { diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/TimeSeriesAuthIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/TimeSeriesAuthIntegrationTest.java index 057317bc2004..41fa6e7ea187 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/TimeSeriesAuthIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/TimeSeriesAuthIntegrationTest.java @@ -49,6 +49,7 @@ protected Map getPinotClientTransportHeaders() { @Override protected void overrideControllerConf(Map properties) { + super.overrideControllerConf(properties); BasicAuthTestUtils.addControllerConfiguration(properties); } diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/TimeSeriesIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/TimeSeriesIntegrationTest.java index ac4f69f3d822..d7490418e162 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/TimeSeriesIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/TimeSeriesIntegrationTest.java @@ -178,6 +178,26 @@ public void testStartTimeEqualsEndTimeQuery() { assertEquals(series.size(), 0); } + @Test + public void testControllerTimeseriesEndpoints() + throws Exception { + // Call /timeseries/api/v1/query_range. + String query = String.format( + "fetch{table=\"mytable_OFFLINE\",filter=\"\",ts_column=\"%s\",ts_unit=\"MILLISECONDS\",value=\"%s\"}" + + " | max{%s} | transformNull{0} | keepLastValue{}", + TS_COLUMN, TOTAL_TRIPS_COLUMN, DEVICE_OS_COLUMN + ); + JsonNode result = getTimeseriesQuery(getControllerBaseApiUrl(), query, QUERY_START_TIME_SEC, QUERY_END_TIME_SEC, + getHeaders()); + assertEquals(result.get("status").asText(), "success"); + + // Call /timeseries/languages. + var statusCodeAndResponse = sendGetRequestWithStatusCode( + getControllerBaseApiUrl() + "/timeseries/languages", getHeaders()); + assertEquals(statusCodeAndResponse.getLeft(), 200); + assertEquals(statusCodeAndResponse.getRight(), "[\"m3ql\"]"); + } + protected Map getHeaders() { return Collections.emptyMap(); } @@ -199,6 +219,15 @@ private void runGroupedTimeSeriesQuery(String query, int expectedGroups, TimeSer } } + @Override + protected void overrideControllerConf(Map properties) { + properties.put(PinotTimeSeriesConfiguration.getEnabledLanguagesConfigKey(), "m3ql"); + properties.put(PinotTimeSeriesConfiguration.getLogicalPlannerConfigKey("m3ql"), + "org.apache.pinot.tsdb.m3ql.M3TimeSeriesPlanner"); + properties.put(PinotTimeSeriesConfiguration.getSeriesBuilderFactoryConfigKey("m3ql"), + SimpleTimeSeriesBuilderFactory.class.getName()); + } + @Override protected void overrideBrokerConf(PinotConfiguration brokerConf) { addTimeSeriesConfigurations(brokerConf); From 43ec504214397467fa030c51ba2fce5e1e21e822 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Sat, 2 Aug 2025 16:50:28 -0700 Subject: [PATCH 063/167] Bump software.amazon.awssdk:bom from 2.32.12 to 2.32.13 (#16483) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 6de6264fdb70..eedc5e4ce296 100644 --- a/pom.xml +++ b/pom.xml @@ -178,7 +178,7 @@ 0.15.1 0.4.7 4.2.2 - 2.32.12 + 2.32.13 1.2.36 1.22.0 2.14.0 From b6f0067e1c0c847c3ee9859261a0876f54f18d59 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Sat, 2 Aug 2025 16:50:43 -0700 Subject: [PATCH 064/167] Bump org.apache.pulsar:pulsar-bom from 4.0.5 to 4.0.6 (#16484) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index eedc5e4ce296..82b9f3514802 100644 --- a/pom.xml +++ b/pom.xml @@ -198,7 +198,7 @@ 2.8.2 3.9.1 7.7.0 - 4.0.5 + 4.0.6 1.20.2 From 9e4bb2a80f82f63b3069de3ef879d270b886021e Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Sat, 2 Aug 2025 21:03:20 -0700 Subject: [PATCH 065/167] Add TableNameExtractor class to improve SQL table name extraction (#16480) - Implement new TableNameExtractor class with robust SQL parsing using Calcite AST - Fix ClassCastException with multi-statement queries (issue #11823) - Support complex SQL constructs: JOINs, CTEs, subqueries, aliases - Use reflection to dynamically load reserved SQL keywords - Add comprehensive test coverage with 100+ test cases - Update Connection and GrpcConnection to use new TableNameExtractor - Remove old resolveTableName method from Connection class - Improve error handling with graceful fallbacks --- .../org/apache/pinot/client/Connection.java | 23 +- .../pinot/client/TableNameExtractor.java | 292 +++++++ .../pinot/client/TableNameExtractorTest.java | 796 ++++++++++++++++++ 3 files changed, 1095 insertions(+), 16 deletions(-) create mode 100644 pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/TableNameExtractor.java create mode 100644 pinot-clients/pinot-java-client/src/test/java/org/apache/pinot/client/TableNameExtractorTest.java diff --git a/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/Connection.java b/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/Connection.java index cc61f4591eae..8ba457981dba 100644 --- a/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/Connection.java +++ b/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/Connection.java @@ -21,12 +21,8 @@ import com.google.common.collect.Iterables; import java.util.List; import java.util.Properties; -import java.util.Set; import java.util.concurrent.CompletableFuture; import javax.annotation.Nullable; -import org.apache.pinot.common.utils.request.RequestUtils; -import org.apache.pinot.sql.parsers.CalciteSqlParser; -import org.apache.pinot.sql.parsers.SqlNodeAndOptions; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -107,8 +103,8 @@ public ResultSetGroup execute(@Nullable String tableName, String query) */ public ResultSetGroup execute(@Nullable Iterable tableNames, String query) throws PinotClientException { - String[] resultTableNames = - (tableNames == null) ? resolveTableName(query) : Iterables.toArray(tableNames, String.class); + String[] resultTableNames = (tableNames == null) ? resolveTableName(query) + : Iterables.toArray(tableNames, String.class); String brokerHostPort = _brokerSelector.selectBroker(resultTableNames); if (brokerHostPort == null) { throw new PinotClientException("Could not find broker to query " + ((tableNames == null) ? "with no tables" @@ -156,8 +152,8 @@ public CompletableFuture executeAsync(@Nullable String tableName */ public CompletableFuture executeAsync(@Nullable Iterable tableNames, String query) throws PinotClientException { - String[] resultTableNames = - (tableNames == null) ? resolveTableName(query) : Iterables.toArray(tableNames, String.class); + String[] resultTableNames = (tableNames == null) ? resolveTableName(query) + : Iterables.toArray(tableNames, String.class); String brokerHostPort = _brokerSelector.selectBroker(resultTableNames); if (brokerHostPort == null) { throw new PinotClientException("Could not find broker to query for statement: " + query); @@ -173,16 +169,11 @@ public CompletableFuture executeAsync(@Nullable Iterable @Nullable public static String[] resolveTableName(String query) { try { - SqlNodeAndOptions sqlNodeAndOptions = CalciteSqlParser.compileToSqlNodeAndOptions(query); - Set tableNames = - RequestUtils.getTableNames(CalciteSqlParser.compileSqlNodeToPinotQuery(sqlNodeAndOptions.getSqlNode())); - if (tableNames != null) { - return tableNames.toArray(new String[0]); - } + return TableNameExtractor.resolveTableName(query); } catch (Exception e) { - LOGGER.error("Cannot parse table name from query: {}. Fallback to broker selector default.", query, e); + LOGGER.warn("Failed to extract table names for query: {}, fall back to default broker selector", query, e); + return null; } - return null; } /** diff --git a/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/TableNameExtractor.java b/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/TableNameExtractor.java new file mode 100644 index 000000000000..26f64a650db0 --- /dev/null +++ b/pinot-clients/pinot-java-client/src/main/java/org/apache/pinot/client/TableNameExtractor.java @@ -0,0 +1,292 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.client; + +import java.util.HashSet; +import java.util.Set; +import javax.annotation.Nullable; +import org.apache.calcite.sql.SqlBasicCall; +import org.apache.calcite.sql.SqlIdentifier; +import org.apache.calcite.sql.SqlJoin; +import org.apache.calcite.sql.SqlKind; +import org.apache.calcite.sql.SqlNode; +import org.apache.calcite.sql.SqlNodeList; +import org.apache.calcite.sql.SqlOrderBy; +import org.apache.calcite.sql.SqlSelect; +import org.apache.calcite.sql.SqlWith; +import org.apache.calcite.sql.SqlWithItem; +import org.apache.pinot.sql.parsers.CalciteSqlParser; +import org.apache.pinot.sql.parsers.SqlCompilationException; +import org.apache.pinot.sql.parsers.SqlNodeAndOptions; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + + +/** + * Helper class to extract table names from Calcite SqlNode tree. + */ +public class TableNameExtractor { + private static final Logger LOGGER = LoggerFactory.getLogger(TableNameExtractor.class); + + /** + * Returns the name of all the tables used in a sql query. + * + * @param query The SQL query string to analyze + * @return name of all the tables used in a sql query, or null if parsing fails + */ + @Nullable + public static String[] resolveTableName(String query) { + try { + SqlNodeAndOptions sqlNodeAndOptions = CalciteSqlParser.compileToSqlNodeAndOptions(query); + Set tableNames = extractTableNamesFromPinotQuery(sqlNodeAndOptions.getSqlNode()); + if (tableNames != null) { + return tableNames.toArray(new String[0]); + } + return null; + } catch (Exception e) { + throw new RuntimeException("Cannot parse table name from query: " + query, e); + } + } + + /** + * Extracts table names from a multi-stage query using Calcite SQL AST traversal. + * + * @param sqlNode The root SqlNode of the parsed query + * @return Set of table names found in the query + */ + private static Set extractTableNamesFromPinotQuery(SqlNode sqlNode) { + TableNameExtractor extractor = new TableNameExtractor(); + extractor.extractTableNames(sqlNode); + return extractor.getTableNames(); + } + + private final Set _tableNames = new HashSet<>(); + private final Set _cteNames = new HashSet<>(); + private boolean _inFromClause = false; + + /** + * Returns the set of table names extracted from the SQL node tree. + *

    + * This method should be called after {@link #extractTableNames(SqlNode)} has been invoked + * to populate the set of table names. + * + * @return Set of table names found in the SQL node tree + */ + public Set getTableNames() { + return _tableNames; + } + + public void extractTableNames(SqlNode node) { + assert node != null; + if (node instanceof SqlWith) { + visitWith((SqlWith) node); + } else if (node instanceof SqlOrderBy) { + visitOrderBy((SqlOrderBy) node); + } else if (node instanceof SqlWithItem) { + visitWithItem((SqlWithItem) node); + } else if (node instanceof SqlSelect) { + visitSelect((SqlSelect) node); + } else if (node instanceof SqlJoin) { + visitJoin((SqlJoin) node); + } else if (node instanceof SqlBasicCall) { + visitBasicCall((SqlBasicCall) node); + } else if (node instanceof SqlIdentifier) { + visitIdentifier((SqlIdentifier) node); + } else if (node instanceof SqlNodeList) { + visitNodeList((SqlNodeList) node); + } + } + + private void visitWith(SqlWith with) { + // Visit the WITH list (CTE definitions) + if (with.withList != null) { + visitNodeList(with.withList); + } + // Visit the main query body + if (with.body != null) { + extractTableNames(with.body); + } + } + + private void visitOrderBy(SqlOrderBy orderBy) { + // Visit the main query - this is the most important part + if (orderBy.query != null) { + extractTableNames(orderBy.query); + } + // Visit ORDER BY expressions for potential subqueries + if (orderBy.orderList != null) { + // Don't set inFromClause=true for ORDER BY expressions + // as they typically contain column references, not table names + visitNodeList(orderBy.orderList); + } + // Visit OFFSET clause if it contains subqueries (rare but possible) + if (orderBy.offset != null) { + extractTableNames(orderBy.offset); + } + // Visit FETCH/LIMIT clause if it contains subqueries (rare but possible) + if (orderBy.fetch != null) { + extractTableNames(orderBy.fetch); + } + } + + private void visitWithItem(SqlWithItem withItem) { + // Track the CTE name so we don't treat it as a table later + if (withItem.name != null) { + String cteName = withItem.name.getSimple(); + _cteNames.add(cteName); + } + // Extract table names from the CTE query definition, not the CTE alias + if (withItem.query != null) { + extractTableNames(withItem.query); + } + } + + private void visitSelect(SqlSelect select) { + // Visit FROM clause - this is where we expect to find table names + if (select.getFrom() != null) { + _inFromClause = true; + extractTableNames(select.getFrom()); + _inFromClause = false; + } + // Visit other clauses for subqueries + if (select.getWhere() != null) { + extractTableNames(select.getWhere()); + } + if (select.getGroup() != null) { + visitNodeList(select.getGroup()); + } + if (select.getHaving() != null) { + extractTableNames(select.getHaving()); + } + if (select.getOrderList() != null) { + visitNodeList(select.getOrderList()); + } + if (select.getSelectList() != null) { + visitNodeList(select.getSelectList()); + } + } + + private void visitJoin(SqlJoin join) { + // Visit both sides of the join - ensure they're processed as FROM clause items + boolean wasInFromClause = _inFromClause; + if (join.getLeft() != null) { + _inFromClause = true; + extractTableNames(join.getLeft()); + } + if (join.getRight() != null) { + _inFromClause = true; + extractTableNames(join.getRight()); + } + // Visit join condition but not as part of FROM clause context + // This handles potential subqueries in join conditions while avoiding + // incorrectly extracting column references as table names + if (join.getCondition() != null) { + _inFromClause = false; + extractTableNames(join.getCondition()); + } + // Restore original context + _inFromClause = wasInFromClause; + } + + private void visitBasicCall(SqlBasicCall call) { + if (call.getKind() == SqlKind.AS) { + // Handle table aliases like "tableA AS a" + // For AS operations, the first operand is the actual table name + if (!call.getOperandList().isEmpty() && call.getOperandList().get(0) != null) { + extractTableNames(call.getOperandList().get(0)); + } + } else if (call.getKind() == SqlKind.WITH) { + // Handle CTE (Common Table Expression) + visitWithClause(call); + } else if (call.getKind() == SqlKind.VALUES) { + // Handle VALUES clause - usually doesn't contain table references + // Skip this to avoid false positives + } else { + // For other basic calls, visit all operands + for (SqlNode operand : call.getOperandList()) { + if (operand != null) { + extractTableNames(operand); + } + } + } + } + + private void visitIdentifier(SqlIdentifier identifier) { + // Only extract table names when we're in a FROM clause + if (_inFromClause && !identifier.names.isEmpty()) { + String tableName = identifier.names.get(identifier.names.size() - 1); + // Filter out system identifiers and CTE names + if (!tableName.startsWith("$") && !_cteNames.contains(tableName)) { + _tableNames.add(tableName); + } + } + } + + /** + * Visit a SqlNodeList by visiting each node in the list. + */ + private void visitNodeList(SqlNodeList nodeList) { + if (nodeList != null) { + for (SqlNode node : nodeList) { + if (node != null) { + extractTableNames(node); + } + } + } + } + + /** + * Handle WITH clause (CTE - Common Table Expression). + */ + private void visitWithClause(SqlNode node) { + try { + // WITH clause typically has operands: [with_list, query] + if (node instanceof SqlBasicCall) { + SqlBasicCall withCall = (SqlBasicCall) node; + for (SqlNode operand : withCall.getOperandList()) { + if (operand != null) { + extractTableNames(operand); + } + } + } + } catch (Exception e) { + // Fallback to generic operand handling + visitNodeOperands(node); + } + } + + /** + * Generic method to visit node operands when specific handling is not available. + */ + private void visitNodeOperands(SqlNode node) { + try { + // Try to access operands through common interface + if (node instanceof SqlBasicCall) { + SqlBasicCall call = (SqlBasicCall) node; + for (SqlNode operand : call.getOperandList()) { + if (operand != null) { + extractTableNames(operand); + } + } + } + } catch (Exception e) { + throw new SqlCompilationException("Exception encountered while visiting node operands: " + node, e); + } + } +} diff --git a/pinot-clients/pinot-java-client/src/test/java/org/apache/pinot/client/TableNameExtractorTest.java b/pinot-clients/pinot-java-client/src/test/java/org/apache/pinot/client/TableNameExtractorTest.java new file mode 100644 index 000000000000..b119e97be5fc --- /dev/null +++ b/pinot-clients/pinot-java-client/src/test/java/org/apache/pinot/client/TableNameExtractorTest.java @@ -0,0 +1,796 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.client; + +import java.util.Arrays; +import java.util.HashSet; +import java.util.Set; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertTrue; + + +/** + * Tests for the TableNameExtractor class. + */ +public class TableNameExtractorTest { + + @Test + public void testResolveTableNameWithSingleQuery() { + // Test that single queries work correctly + String singleQuery = "SELECT * FROM myTable WHERE id > 100"; + + String[] tableNames = TableNameExtractor.resolveTableName(singleQuery); + + assertNotNull(tableNames, "Table names should not be null"); + assertEquals(tableNames.length, 1, "Should resolve exactly one table"); + assertEquals(tableNames[0], "myTable", "Should resolve the correct table name"); + } + + @Test + public void testResolveTableNameWithSingleStatementAlias() { + String singleStatementQuery = "SELECT stats.* FROM airlineStats stats LIMIT 10"; + String[] tableNames = TableNameExtractor.resolveTableName(singleStatementQuery); + + assertNotNull(tableNames); + assertEquals(tableNames.length, 1); + assertEquals(tableNames[0], "airlineStats"); + } + + @Test + public void testResolveTableNameWithMultiStatementQuery() { + // Test the fix for issue #11823: CalciteSQLParser error with multi-statement queries + String multiStatementQuery = "SET useMultistageEngine=true;\nSELECT stats.* FROM airlineStats stats LIMIT 10"; + + // This should not throw a ClassCastException anymore + String[] tableNames = TableNameExtractor.resolveTableName(multiStatementQuery); + + // Should successfully resolve the table name + assertNotNull(tableNames, "Table names should not be null"); + assertEquals(tableNames.length, 1, "Should resolve exactly one table"); + assertEquals(tableNames[0], "airlineStats", "Should resolve the correct table name"); + } + + @Test + public void testResolveTableNameWithMultipleSetStatements() { + // Test with multiple SET statements + String multiSetQuery = "SET useMultistageEngine=true;\nSET timeoutMs=10000;\nSELECT * FROM testTable"; + + String[] tableNames = TableNameExtractor.resolveTableName(multiSetQuery); + + assertNotNull(tableNames, "Table names should not be null"); + assertEquals(tableNames.length, 1, "Should resolve exactly one table"); + assertEquals(tableNames[0], "testTable", "Should resolve the correct table name"); + } + + @Test + public void testResolveTableNameWithMultipleSetStatementsAndJoin() { + String multiStatementQuery = "SET useMultistageEngine=true;\nSET maxRowsInJoin=1000;\n" + + "SELECT stats.* FROM airlineStats stats LIMIT 10"; + String[] tableNames = TableNameExtractor.resolveTableName(multiStatementQuery); + + assertNotNull(tableNames, "Table names should be resolved for queries with multiple SET statements"); + assertEquals(tableNames.length, 1); + assertEquals(tableNames[0], "airlineStats"); + } + + @Test + public void testResolveTableNameWithJoin() { + // Test with JOIN queries + String joinQuery = "SELECT * FROM table1 t1 JOIN table2 t2 ON t1.id = t2.id"; + + String[] tableNames = TableNameExtractor.resolveTableName(joinQuery); + + assertNotNull(tableNames, "Table names should not be null"); + assertEquals(tableNames.length, 2, "Should resolve two tables"); + assertTrue(Arrays.asList(tableNames).contains("table1"), "Should contain table1"); + assertTrue(Arrays.asList(tableNames).contains("table2"), "Should contain table2"); + } + + @Test + public void testResolveTableNameWithJoinQueryAndSetStatements() { + String joinQuery = "SET useMultistageEngine=true;\n" + + "SELECT a.col1, b.col2 FROM tableA a JOIN tableB b ON a.id = b.id"; + String[] tableNames = TableNameExtractor.resolveTableName(joinQuery); + + assertNotNull(tableNames, "Table names should be resolved for join queries with SET statements"); + assertEquals(tableNames.length, 2); + + Set expectedTableNames = new HashSet<>(Arrays.asList("tableA", "tableB")); + Set actualTableNames = new HashSet<>(Arrays.asList(tableNames)); + assertEquals(actualTableNames, expectedTableNames); + } + + @Test + public void testResolveTableNameWithExplicitAlias() { + // Test with explicit AS alias + String aliasQuery = "SELECT u.name FROM users AS u WHERE u.active = true"; + + String[] tableNames = TableNameExtractor.resolveTableName(aliasQuery); + + assertNotNull(tableNames, "Table names should not be null"); + assertEquals(tableNames.length, 1, "Should resolve exactly one table"); + assertEquals(tableNames[0], "users", "Should resolve the actual table name, not the alias"); + } + + @Test + public void testResolveTableNameWithImplicitAlias() { + // Test with implicit alias (no AS keyword) + String implicitAliasQuery = "SELECT o.id, u.name FROM orders o JOIN users u ON o.user_id = u.id"; + + String[] tableNames = TableNameExtractor.resolveTableName(implicitAliasQuery); + + assertNotNull(tableNames, "Table names should not be null"); + assertEquals(tableNames.length, 2, "Should resolve two tables"); + assertTrue(Arrays.asList(tableNames).contains("orders"), "Should contain orders table"); + assertTrue(Arrays.asList(tableNames).contains("users"), "Should contain users table"); + } + + @Test + public void testResolveTableNameWithCTE() { + // Test with Common Table Expression (CTE) + String cteQuery = "WITH active_users AS (SELECT * FROM users WHERE active = true) " + + "SELECT au.name FROM active_users au JOIN orders o ON au.id = o.user_id"; + + String[] tableNames = TableNameExtractor.resolveTableName(cteQuery); + + assertNotNull(tableNames, "Table names should not be null"); + assertEquals(tableNames.length, 2, "Should resolve two tables"); + assertTrue(Arrays.asList(tableNames).contains("users"), "Should contain users table from CTE"); + assertTrue(Arrays.asList(tableNames).contains("orders"), "Should contain orders table"); + } + + @Test + public void testResolveTableNameWithNestedCTE() { + // Test with nested CTEs + String nestedCteQuery = "WITH user_orders AS (" + + " SELECT u.id, u.name, o.order_date " + + " FROM users u JOIN orders o ON u.id = o.user_id" + + "), recent_orders AS (" + + " SELECT * FROM user_orders WHERE order_date > '2023-01-01'" + + ") " + + "SELECT ro.name FROM recent_orders ro JOIN products p ON ro.id = p.user_id"; + + String[] tableNames = TableNameExtractor.resolveTableName(nestedCteQuery); + + assertNotNull(tableNames, "Table names should not be null"); + assertEquals(tableNames.length, 3, "Should resolve three tables"); + assertTrue(Arrays.asList(tableNames).contains("users"), "Should contain users table"); + assertTrue(Arrays.asList(tableNames).contains("orders"), "Should contain orders table"); + assertTrue(Arrays.asList(tableNames).contains("products"), "Should contain products table"); + } + + @Test + public void testResolveTableNameWithSubqueryAlias() { + // Test with subquery alias + String subqueryQuery = "SELECT t.name FROM (SELECT * FROM users WHERE active = true) AS t " + + "JOIN orders o ON t.id = o.user_id"; + + String[] tableNames = TableNameExtractor.resolveTableName(subqueryQuery); + + assertNotNull(tableNames, "Table names should not be null"); + assertEquals(tableNames.length, 2, "Should resolve two tables"); + assertTrue(Arrays.asList(tableNames).contains("users"), "Should contain users table from subquery"); + assertTrue(Arrays.asList(tableNames).contains("orders"), "Should contain orders table"); + } + + @Test + public void testResolveTableNameWithComplexJoinAndAliases() { + // Test with multiple JOINs and various alias styles + String complexQuery = "SELECT u.name, o.total, p.title " + + "FROM users AS u " + + "INNER JOIN orders o ON u.id = o.user_id " + + "LEFT JOIN order_items oi ON o.id = oi.order_id " + + "RIGHT JOIN products AS p ON oi.product_id = p.id " + + "WHERE u.active = true"; + + String[] tableNames = TableNameExtractor.resolveTableName(complexQuery); + + assertNotNull(tableNames, "Table names should not be null"); + assertEquals(tableNames.length, 4, "Should resolve four tables"); + assertTrue(Arrays.asList(tableNames).contains("users"), "Should contain users table"); + assertTrue(Arrays.asList(tableNames).contains("orders"), "Should contain orders table"); + assertTrue(Arrays.asList(tableNames).contains("order_items"), "Should contain order_items table"); + assertTrue(Arrays.asList(tableNames).contains("products"), "Should contain products table"); + } + + @Test + public void testResolveTableNameWithJoinConditionSubquery() { + // Test with subquery in join condition + String joinSubqueryQuery = "SELECT u.name, o.total " + + "FROM users u " + + "JOIN orders o ON u.id = o.user_id " + + "AND o.id IN (SELECT order_id FROM order_items WHERE quantity > 5)"; + + String[] tableNames = TableNameExtractor.resolveTableName(joinSubqueryQuery); + + assertNotNull(tableNames, "Table names should not be null"); + assertEquals(tableNames.length, 3, "Should resolve three tables"); + assertTrue(Arrays.asList(tableNames).contains("users"), "Should contain users table"); + assertTrue(Arrays.asList(tableNames).contains("orders"), "Should contain orders table"); + assertTrue(Arrays.asList(tableNames).contains("order_items"), + "Should contain order_items table from subquery"); + } + + @Test + public void testResolveTableNameWithOrderBy() { + // Test with ORDER BY clause + String orderByQuery = "SELECT * FROM users ORDER BY name"; + String[] tableNames = TableNameExtractor.resolveTableName(orderByQuery); + + assertNotNull(tableNames, "Table names should not be null"); + assertEquals(tableNames.length, 1, "Should resolve exactly one table"); + assertEquals(tableNames[0], "users", "Should resolve the correct table name"); + } + + @Test + public void testResolveTableNameWithOrderBySubquery() { + // Test with subquery in ORDER BY clause (rare but possible) + String orderBySubqueryQuery = "SELECT * FROM users u ORDER BY " + + "(SELECT COUNT(*) FROM orders o WHERE o.user_id = u.id)"; + String[] tableNames = TableNameExtractor.resolveTableName(orderBySubqueryQuery); + + assertNotNull(tableNames, "Table names should not be null"); + assertEquals(tableNames.length, 2, "Should resolve two tables"); + assertTrue(Arrays.asList(tableNames).contains("users"), "Should contain users table"); + assertTrue(Arrays.asList(tableNames).contains("orders"), "Should contain orders table"); + } + + @Test(expectedExceptions = RuntimeException.class) + public void testResolveTableNameWithInvalidQuery() { + String[] tableNames = TableNameExtractor.resolveTableName("INVALID SQL QUERY"); + } + + @Test(expectedExceptions = RuntimeException.class) + public void testResolveTableNameWithOnlySetStatements() { + TableNameExtractor.resolveTableName("SET useMultistageEngine=true;"); + } + + @Test(expectedExceptions = RuntimeException.class) + public void testResolveTableNameWithNullQuery() { + TableNameExtractor.resolveTableName(null); + } + + @Test(expectedExceptions = RuntimeException.class) + public void testResolveTableNameWithEmptyQuery() { + TableNameExtractor.resolveTableName(""); + } + + /** + * Data provider for SQL queries and their expected table names. + * This makes it easy to add new test cases by simply adding entries to this array. + * + * @return Object[][] where each Object[] contains: [testName (String), sqlQuery (String), + * expectedTableNames (String[] or null)] + * Each entry in the returned array is an Object[] of length 3, structured as follows: + *

      + *
    • testName (String): A descriptive name for the test case.
      • + *
    • sqlQuery (String): The SQL query to be tested.
      • + *
    • expectedTableNames (String[]): The expected table names to be extracted from the query, + * or {@code null} if no table names are expected (e.g., for invalid or empty queries).
      • + *
    + * This makes it easy to add new test cases by simply adding entries to this array. + */ + @DataProvider(name = "sqlQueries") + public Object[][] sqlQueriesDataProvider() { + return new Object[][]{ + // Basic queries + { + "Simple SELECT", + "SELECT * FROM users", + new String[]{"users"}, + false + }, + { + "SELECT with WHERE", + "SELECT name FROM users WHERE age > 18", + new String[]{"users"}, + false + }, + { + "SELECT with LIMIT", + "SELECT * FROM products LIMIT 10", + new String[]{"products"}, + false + }, + + // Aliases + { + "Explicit alias", + "SELECT u.name FROM users u", + new String[]{"users"}, + false + }, + { + "Implicit alias", + "SELECT u.name FROM users u JOIN orders o ON u.id = o.user_id", + new String[]{"users", "orders"}, + false + }, + { + "Multiple aliases", + "SELECT t1.col1, t2.col2 FROM table1 t1, table2 t2", + new String[]{"table1", "table2"}, + false + }, + + // JOINs + { + "INNER JOIN", + "SELECT * FROM users u INNER JOIN orders o ON u.id = o.user_id", + new String[]{"users", "orders"}, + false + }, + { + "LEFT JOIN", + "SELECT * FROM users u LEFT JOIN orders o ON u.id = o.user_id", + new String[]{"users", "orders"}, + false + }, + { + "RIGHT JOIN", + "SELECT * FROM users u RIGHT JOIN orders o ON u.id = o.user_id", + new String[]{"users", "orders"}, + false + }, + { + "FULL JOIN", + "SELECT * FROM users u FULL JOIN orders o ON u.id = o.user_id", + new String[]{"users", "orders"}, + false + }, + { + "Multiple JOINs", + "SELECT * FROM users u JOIN orders o ON u.id = o.user_id JOIN products p ON o.product_id = p.id", + new String[]{"users", "orders", "products"}, + false + }, + + // CTEs (Common Table Expressions) + { + "Simple CTE", + "WITH active_users AS (SELECT * FROM users WHERE active = true) " + + "SELECT * FROM active_users", + new String[]{"users"}, + false + }, + { + "Multiple CTEs", + "WITH active_users AS (SELECT * FROM users WHERE active = true), " + + "recent_orders AS (SELECT * FROM orders WHERE created_date > '2024-01-01') " + + "SELECT au.name, ro.order_id FROM active_users au JOIN recent_orders ro ON au.id = ro.user_id", + new String[]{"users", "orders"}, + false + }, + { + "Nested CTE", + "WITH user_stats AS (SELECT user_id, COUNT(*) as order_count FROM orders GROUP BY user_id), " + + "top_users AS (SELECT * FROM user_stats WHERE order_count > 10) " + + "SELECT u.name, tu.order_count FROM users u JOIN top_users tu ON u.id = tu.user_id", + new String[]{"orders", "users"}, + false + }, + + // Subqueries + { + "Subquery in FROM", + "SELECT * FROM (SELECT * FROM users WHERE active = true) AS active_users", + new String[]{"users"}, + false + }, + { + "Subquery in JOIN", + "SELECT u.name FROM users u JOIN (SELECT user_id FROM orders WHERE amount > 100) o ON u.id = o.user_id", + new String[]{"users", "orders"}, + false + }, + { + "Subquery in WHERE", + "SELECT * FROM users WHERE id IN (SELECT user_id FROM orders WHERE amount > 100)", + new String[]{"users", "orders"}, + false + }, + + // Multi-statement queries + { + "SET + SELECT", + "SET useMultistageEngine=true; SELECT * FROM users", + new String[]{"users"}, + false + }, + { + "Multiple SETs", + "SET useMultistageEngine=true; SET timeoutMs=10000; SELECT * FROM products", + new String[]{"products"}, + false + }, + { + "SET + JOIN", + "SET useMultistageEngine=true; SELECT u.name FROM users u JOIN orders o ON u.id = o.user_id", + new String[]{"users", "orders"}, + false + }, + + // Complex queries + { + "Complex query with all features", + "SET useMultistageEngine=true; " + + "WITH user_stats AS (SELECT user_id, COUNT(*) as order_count FROM orders GROUP BY user_id) " + + "SELECT u.name, us.order_count " + + "FROM users u " + + "JOIN user_stats us ON u.id = us.user_id " + + "JOIN (SELECT user_id FROM products WHERE category = 'electronics') p ON u.id = p.user_id " + + "WHERE us.order_count > 5 " + + "ORDER BY us.order_count DESC", + new String[]{"orders", "users", "products"}, + false + }, + + // Edge cases + { + "Table with underscore", + "SELECT * FROM user_profiles", + new String[]{"user_profiles"}, + false + }, + { + "Table with numbers", + "SELECT * FROM table_2024", + new String[]{"table_2024"}, + false + }, + { + "Multiple tables same name", + "SELECT * FROM users u1 JOIN users u2 ON u1.id = u2.referrer_id", + new String[]{"users"}, + false + }, + + // Queries that should throw exception + { + "Only SET statements", + "SET useMultistageEngine=true; SET timeoutMs=10000;", + null, + true + }, + { + "Empty query", + "", + null, + true + }, + { + "Null query", + null, + null, + true + }, + { + "Invalid SQL", + "INVALID SQL QUERY", + null, + true + }, + + // Additional queries from BaseClusterIntegrationTestSet + // Basic aggregation queries + { + "SUM INTEGER", + "SELECT SUM(ActualElapsedTime) FROM mytable", + new String[]{"mytable"}, + false + }, + { + "SUM FLOAT", + "SELECT SUM(CAST(ActualElapsedTime AS FLOAT)) FROM mytable", + new String[]{"mytable"}, + false + }, + { + "SUM DOUBLE", + "SELECT SUM(CAST(ActualElapsedTime AS DOUBLE)) FROM mytable", + new String[]{"mytable"}, + false + }, + { + "COUNT with WHERE", + "SELECT COUNT(*) FROM mytable WHERE CarrierDelay=15 AND ArrDelay > CarrierDelay LIMIT 1", + new String[]{"mytable"}, + false + }, + { + "MAX MIN", + "SELECT MAX(Quarter), MAX(FlightNum) FROM mytable LIMIT 8", + new String[]{"mytable"}, + false + }, + + // Complex SELECT with arithmetic and functions + { + "Arithmetic in SELECT", + "SELECT ArrDelay, CarrierDelay, (ArrDelay - CarrierDelay) AS diff, " + + "substring(DestStateName, 4, 8) as stateSubStr FROM mytable WHERE CarrierDelay=15 AND " + + "ArrDelay > CarrierDelay ORDER BY diff, ArrDelay, CarrierDelay LIMIT 100000", + new String[]{"mytable"}, + false + }, + { + "Arithmetic operations", + "SELECT ArrTime, ArrTime * 10 FROM mytable WHERE DaysSinceEpoch >= 16312", + new String[]{"mytable"}, + false + }, + { + "Complex arithmetic", + "SELECT ArrTime, ArrTime + ArrTime * 9 - ArrTime * 10 FROM mytable WHERE DaysSinceEpoch >= 16312", + new String[]{"mytable"}, + false + }, + + // GROUP BY queries + { + "GROUP BY with aggregation", + "SELECT COUNT(*), MAX(ArrTime), MIN(ArrTime), DaysSinceEpoch FROM mytable GROUP BY DaysSinceEpoch", + new String[]{"mytable"}, + false + }, + { + "GROUP BY with ORDER BY", + "SELECT DaysSinceEpoch, COUNT(*), MAX(ArrTime), MIN(ArrTime) FROM mytable GROUP BY DaysSinceEpoch", + new String[]{"mytable"}, + false + }, + + // HAVING clauses + { + "HAVING clause", + "SELECT COUNT(*) AS Count, DaysSinceEpoch FROM mytable GROUP BY DaysSinceEpoch HAVING Count > 350", + new String[]{"mytable"}, + false + }, + { + "HAVING with arithmetic", + "SELECT MAX(ArrDelay) - MAX(AirTime) AS Diff, DaysSinceEpoch FROM mytable " + + "GROUP BY DaysSinceEpoch HAVING Diff * 2 > 1000 ORDER BY Diff ASC", + new String[]{"mytable"}, + false + }, + + // LIKE patterns + { + "LIKE pattern", + "SELECT count(*) FROM mytable WHERE OriginState LIKE 'A_'", + new String[]{"mytable"}, + false + }, + { + "LIKE with %", + "SELECT count(*) FROM mytable WHERE DestCityName LIKE 'C%'", + new String[]{"mytable"}, + false + }, + { + "LIKE with _ and %", + "SELECT count(*) FROM mytable WHERE DestCityName LIKE '_h%'", + new String[]{"mytable"}, + false + }, + + // NOT operators + { + "NOT BETWEEN", + "SELECT count(*) FROM mytable WHERE OriginState NOT BETWEEN 'DE' AND 'PA'", + new String[]{"mytable"}, + false + }, + { + "NOT LIKE", + "SELECT count(*) FROM mytable WHERE OriginState NOT LIKE 'A_'", + new String[]{"mytable"}, + false + }, + { + "NOT with parentheses", + "SELECT count(*) FROM mytable WHERE NOT (DaysSinceEpoch = 16312 AND Carrier = 'DL')", + new String[]{"mytable"}, + false + }, + + // CAST operations + { + "CAST operations", + "SELECT SUM(CAST(CAST(ArrTime AS VARCHAR) AS LONG)) FROM mytable " + + "WHERE DaysSinceEpoch <> 16312 AND Carrier = 'DL'", + new String[]{"mytable"}, + false + }, + { + "CAST with ORDER BY", + "SELECT CAST(CAST(ArrTime AS STRING) AS BIGINT) FROM mytable " + + "WHERE DaysSinceEpoch <> 16312 AND Carrier = 'DL' ORDER BY ArrTime DESC", + new String[]{"mytable"}, + false + }, + + // DateTime functions + { + "DateTimeConvert", + "SELECT dateTimeConvert(DaysSinceEpoch,'1:DAYS:EPOCH','1:HOURS:EPOCH','1:HOURS'), COUNT(*) FROM mytable " + + "GROUP BY dateTimeConvert(DaysSinceEpoch,'1:DAYS:EPOCH','1:HOURS:EPOCH','1:HOURS') " + + "ORDER BY COUNT(*), dateTimeConvert(DaysSinceEpoch,'1:DAYS:EPOCH','1:HOURS:EPOCH','1:HOURS') DESC", + new String[]{"mytable"}, + false + }, + { + "TimeConvert", + "SELECT timeConvert(DaysSinceEpoch,'DAYS','SECONDS'), COUNT(*) FROM mytable " + + "GROUP BY timeConvert(DaysSinceEpoch,'DAYS','SECONDS') " + + "ORDER BY COUNT(*), timeConvert(DaysSinceEpoch,'DAYS','SECONDS') DESC", + new String[]{"mytable"}, + false + }, + + // CASE WHEN statements + { + "CASE WHEN with aggregation", + "SELECT AirlineID, " + + "CASE WHEN Sum(ArrDelay) < 0 THEN 0 WHEN SUM(ArrDelay) > 0 THEN SUM(ArrDelay) END AS SumArrDelay " + + "FROM mytable GROUP BY AirlineID", + new String[]{"mytable"}, + false + }, + { + "CASE WHEN without GROUP BY", + "SELECT CASE WHEN Sum(ArrDelay) < 0 THEN 0 WHEN SUM(ArrDelay) > 0 THEN SUM(ArrDelay) END AS SumArrDelay " + + "FROM mytable", + new String[]{"mytable"}, + false + }, + + // Post-aggregation operations + { + "Post-aggregation in ORDER BY", + "SELECT MAX(ArrTime) FROM mytable GROUP BY DaysSinceEpoch ORDER BY MAX(ArrTime) - MIN(ArrTime)", + new String[]{"mytable"}, + false + }, + { + "Post-aggregation in SELECT", + "SELECT MAX(ArrDelay) + MAX(AirTime) FROM mytable", + new String[]{"mytable"}, + false + }, + + // Virtual columns (these should be treated as regular columns for table name extraction) + { + "Virtual columns", + "SELECT $docId, $segmentName, $hostName FROM mytable", + new String[]{"mytable"}, + false + }, + { + "Virtual columns with WHERE", + "SELECT $docId, $segmentName, $hostName FROM mytable WHERE $docId < 5 LIMIT 50", + new String[]{"mytable"}, + false + }, + { + "Virtual columns with GROUP BY", + "SELECT max($docId) FROM mytable GROUP BY $segmentName", + new String[]{"mytable"}, + false + }, + + // Complex WHERE conditions + { + "Complex WHERE with multiple conditions", + "SELECT count(*) FROM mytable WHERE AirlineID > 20355 AND " + + "OriginState BETWEEN 'PA' AND 'DE' AND DepTime <> 2202 LIMIT 21", + new String[]{"mytable"}, + false + }, + { + "WHERE with arithmetic", + "SELECT ArrTime, ArrTime + ArrTime * 9 - ArrTime * 10 FROM mytable WHERE ArrTime - 100 > 0", + new String[]{"mytable"}, + false + }, + + // Subquery patterns (from V2 tests) + { + "IN_SUBQUERY", + "SELECT COUNT(*) FROM mytable WHERE INSUBQUERY(DestAirportID, " + + "'SELECT IDSET(DestAirportID) FROM mytable WHERE DaysSinceEpoch = 16430') = 1", + new String[]{"mytable"}, + false + }, + { + "NOT IN_SUBQUERY", + "SELECT COUNT(*) FROM mytable WHERE INSUBQUERY(DestAirportID, " + + "'SELECT IDSET(DestAirportID) FROM mytable WHERE DaysSinceEpoch = 16430') = 0", + new String[]{"mytable"}, + false + }, + + // Multi-value column queries + { + "Multi-value IN", + "SELECT DistanceGroup FROM mytable WHERE \"Month\" BETWEEN 1 AND 1 AND " + + "arrayToMV(DivAirportSeqIDs) IN (1078102, 1142303, 1530402, 1172102, 1291503) OR SecurityDelay IN " + + "(1, 0, 14, -9999) LIMIT 10", + new String[]{"mytable"}, + false + }, + + // Options and hints + { + "Query with options", + "SELECT count(*) FROM mytable WHERE OriginState LIKE 'A_' option(orderedPreferredPools=0|1)", + new String[]{"mytable"}, + false + }, + { + "SET with query", + "SET orderedPreferredPools='0 | 1'; SELECT count(*) FROM mytable WHERE OriginState LIKE 'A_'", + new String[]{"mytable"}, + false + } + }; + } + + /** + * Test method that uses the DataProvider to test multiple SQL queries. + * This makes it easy to add new test cases by simply adding entries to the data provider. + * + * @param testName The name of the test case for better reporting + * @param sqlQuery The SQL query to test + * @param expectedTableNames The expected table names that should be extracted + */ + @Test(dataProvider = "sqlQueries") + public void testResolveTableNameWithDataProvider(String testName, String sqlQuery, String[] expectedTableNames, + boolean throwException) { + try { + // Extract table names from the SQL query + String[] actualTableNames = TableNameExtractor.resolveTableName(sqlQuery); + + if (expectedTableNames == null) { + // For queries that should return null (invalid, empty, etc.) + assertNull(actualTableNames, "Query should return null: " + testName); + } else { + // For valid queries, check that we got the expected table names + assertNotNull(actualTableNames, "Table names should not be null for: " + testName); + assertEquals(actualTableNames.length, expectedTableNames.length, + "Should extract correct number of tables for: " + testName); + + // Convert arrays to sets for order-independent comparison + Set actualSet = new HashSet<>(Arrays.asList(actualTableNames)); + Set expectedSet = new HashSet<>(Arrays.asList(expectedTableNames)); + + assertEquals(actualSet, expectedSet, + "Should extract correct table names for: " + testName); + } + } catch (Exception e) { + assertTrue(throwException); + } + } +} From a09d01faae93e5156884d9429decdd43aa8f5582 Mon Sep 17 00:00:00 2001 From: Rajat Venkatesh <1638298+vrajat@users.noreply.github.com> Date: Mon, 4 Aug 2025 15:57:15 +0530 Subject: [PATCH 066/167] Remove unnecessary methods and config for ThreadResourceUsageAccountant (#16490) --- .../PerQueryCPUMemAccountantFactory.java | 20 --------- .../ResourceUsageAccountantFactory.java | 16 ------- .../runtime/operator/MultiStageOperator.java | 2 +- .../operator/MultiStageAccountingTest.java | 2 +- ...MultistageResourceUsageAccountingTest.java | 2 +- .../queries/PerQueryCPUMemAccountantTest.java | 43 ------------------- .../ThreadResourceUsageAccountant.java | 5 --- .../org/apache/pinot/spi/trace/Tracing.java | 13 ------ .../pinot/spi/utils/CommonConstants.java | 3 -- ...rottleOnCriticalHeapUsageExecutorTest.java | 4 -- 10 files changed, 3 insertions(+), 107 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java index b88f06f2917d..293edff6ff68 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/accounting/PerQueryCPUMemAccountantFactory.java @@ -127,9 +127,6 @@ public static class PerQueryCPUMemResourceUsageAccountant implements ThreadResou // track memory usage protected final boolean _isThreadMemorySamplingEnabled; - // is sampling allowed for MSE queries - protected final boolean _isThreadSamplingEnabledForMSE; - protected final Set _inactiveQuery; protected Set _cancelSentQueries; @@ -148,7 +145,6 @@ protected PerQueryCPUMemResourceUsageAccountant(PinotConfiguration config, boole _config = config; _isThreadCPUSamplingEnabled = isThreadCPUSamplingEnabled; _isThreadMemorySamplingEnabled = isThreadMemorySamplingEnabled; - _isThreadSamplingEnabledForMSE = isThreadSamplingEnabledForMSE; _inactiveQuery = inactiveQuery; _instanceId = instanceId; _instanceType = instanceType; @@ -184,11 +180,6 @@ public PerQueryCPUMemResourceUsageAccountant(PinotConfiguration config, String i LOGGER.info("_isThreadCPUSamplingEnabled: {}, _isThreadMemorySamplingEnabled: {}", _isThreadCPUSamplingEnabled, _isThreadMemorySamplingEnabled); - _isThreadSamplingEnabledForMSE = - config.getProperty(CommonConstants.Accounting.CONFIG_OF_ENABLE_THREAD_SAMPLING_MSE, - CommonConstants.Accounting.DEFAULT_ENABLE_THREAD_SAMPLING_MSE); - LOGGER.info("_isThreadSamplingEnabledForMSE: {}", _isThreadSamplingEnabledForMSE); - _queryCancelCallbacks = CacheBuilder.newBuilder().maximumSize( config.getProperty(CommonConstants.Accounting.CONFIG_OF_CANCEL_CALLBACK_CACHE_MAX_SIZE, CommonConstants.Accounting.DEFAULT_CANCEL_CALLBACK_CACHE_MAX_SIZE)).expireAfterWrite( @@ -274,17 +265,6 @@ public void sampleUsage() { sampleThreadCPUTime(); } - /** - * Sample Usage for Multi-stage engine queries - */ - @Override - public void sampleUsageMSE() { - if (_isThreadSamplingEnabledForMSE) { - sampleThreadBytesAllocated(); - sampleThreadCPUTime(); - } - } - @Override public boolean throttleQuerySubmission() { return getWatcherTask().getHeapUsageBytes() > getWatcherTask().getQueryMonitorConfig().getAlarmingLevel(); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/accounting/ResourceUsageAccountantFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/accounting/ResourceUsageAccountantFactory.java index e00f74d888d6..93c68ca7fe38 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/accounting/ResourceUsageAccountantFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/accounting/ResourceUsageAccountantFactory.java @@ -82,9 +82,6 @@ public static class ResourceUsageAccountant implements ThreadResourceUsageAccoun // track memory usage private final boolean _isThreadMemorySamplingEnabled; - // is sampling allowed for MSE queries - private final boolean _isThreadSamplingEnabledForMSE; - private final WatcherTask _watcherTask; private final EnumMap _resourceAggregators; @@ -109,11 +106,6 @@ public ResourceUsageAccountant(PinotConfiguration config, String instanceId, Ins LOGGER.info("_isThreadCPUSamplingEnabled: {}, _isThreadMemorySamplingEnabled: {}", _isThreadCPUSamplingEnabled, _isThreadMemorySamplingEnabled); - _isThreadSamplingEnabledForMSE = - config.getProperty(CommonConstants.Accounting.CONFIG_OF_ENABLE_THREAD_SAMPLING_MSE, - CommonConstants.Accounting.DEFAULT_ENABLE_THREAD_SAMPLING_MSE); - LOGGER.info("_isThreadSamplingEnabledForMSE: {}", _isThreadSamplingEnabledForMSE); - _watcherTask = new WatcherTask(); _resourceAggregators = new EnumMap<>(TrackingScope.class); @@ -138,14 +130,6 @@ public void sampleUsage() { sampleThreadCPUTime(); } - @Override - public void sampleUsageMSE() { - if (_isThreadSamplingEnabledForMSE) { - sampleThreadBytesAllocated(); - sampleThreadCPUTime(); - } - } - @Override public boolean isAnchorThreadInterrupted() { ThreadExecutionContext context = _threadLocalEntry.get().getCurrentThreadTaskStatus(); diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultiStageOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultiStageOperator.java index 3b179ed4b807..adf02c497f40 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultiStageOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/MultiStageOperator.java @@ -97,7 +97,7 @@ protected void sampleAndCheckInterruption(long deadlineMs) { earlyTerminate(); throw QueryErrorCode.EXECUTION_TIMEOUT.asException("Timing out on " + getExplainName()); } - Tracing.ThreadAccountantOps.sampleMSE(); + Tracing.ThreadAccountantOps.sample(); if (Tracing.ThreadAccountantOps.isInterrupted()) { earlyTerminate(); throw QueryErrorCode.SERVER_RESOURCE_LIMIT_EXCEEDED.asException("Resource limit exceeded for operator: " diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MultiStageAccountingTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MultiStageAccountingTest.java index de34a977d05c..57b27b907fdc 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MultiStageAccountingTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MultiStageAccountingTest.java @@ -92,7 +92,7 @@ public static void setUpClass() { configs.put(CommonConstants.Accounting.CONFIG_OF_OOM_PROTECTION_KILLING_QUERY, true); // init accountant and start watcher task Tracing.unregisterThreadAccountant(); - Tracing.ThreadAccountantOps.initializeThreadAccountant(new PinotConfiguration(configs), "testGroupBy", + Tracing.ThreadAccountantOps.createThreadAccountant(new PinotConfiguration(configs), "testGroupBy", InstanceType.SERVER); Tracing.ThreadAccountantOps.startThreadAccountant(); diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MultistageResourceUsageAccountingTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MultistageResourceUsageAccountingTest.java index 534f1f413d54..8312bd098bac 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MultistageResourceUsageAccountingTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/MultistageResourceUsageAccountingTest.java @@ -92,7 +92,7 @@ public static void setUpClass() { // init accountant and start watcher task PinotConfiguration pinotCfg = new PinotConfiguration(configs); Tracing.unregisterThreadAccountant(); - Tracing.ThreadAccountantOps.initializeThreadAccountant(pinotCfg, "testGroupBy", InstanceType.SERVER); + Tracing.ThreadAccountantOps.createThreadAccountant(pinotCfg, "testGroupBy", InstanceType.SERVER); Tracing.ThreadAccountantOps.startThreadAccountant(); // Setup Thread Context diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/PerQueryCPUMemAccountantTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/PerQueryCPUMemAccountantTest.java index 42915480a9a0..1aa4b23af297 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/PerQueryCPUMemAccountantTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/queries/PerQueryCPUMemAccountantTest.java @@ -27,7 +27,6 @@ import java.util.concurrent.Executors; import org.apache.pinot.common.response.broker.ResultTable; import org.apache.pinot.core.accounting.PerQueryCPUMemAccountantFactory; -import org.apache.pinot.core.accounting.ResourceUsageAccountantFactory; import org.apache.pinot.query.planner.physical.DispatchableSubPlan; import org.apache.pinot.spi.accounting.QueryResourceTracker; import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant; @@ -98,48 +97,6 @@ void testWithPerQueryAccountantFactory() { } } - @Test - void testDisableSamplingForMSE() { - HashMap configs = getAccountingConfig(); - configs.put(CommonConstants.Accounting.CONFIG_OF_ENABLE_THREAD_SAMPLING_MSE, false); - - ThreadResourceUsageProvider.setThreadMemoryMeasurementEnabled(true); - PerQueryCPUMemAccountantFactory.PerQueryCPUMemResourceUsageAccountant accountant = - new PerQueryCPUMemAccountantFactory.PerQueryCPUMemResourceUsageAccountant(new PinotConfiguration(configs), - "testWithPerQueryAccountantFactory", InstanceType.SERVER); - - try (MockedStatic tracing = Mockito.mockStatic(Tracing.class, Mockito.CALLS_REAL_METHODS)) { - tracing.when(Tracing::getThreadAccountant).thenReturn(accountant); - ResultTable resultTable = queryRunner("SELECT * FROM a LIMIT 2", false).getResultTable(); - Assert.assertEquals(resultTable.getRows().size(), 2); - - Map resources = accountant.getQueryResources(); - Assert.assertEquals(resources.size(), 1); - Assert.assertEquals(resources.entrySet().iterator().next().getValue().getAllocatedBytes(), 0); - } - } - - @Test - void testDisableSamplingWithResourceUsageAccountantForMSE() { - HashMap configs = getAccountingConfig(); - configs.put(CommonConstants.Accounting.CONFIG_OF_ENABLE_THREAD_SAMPLING_MSE, false); - - ThreadResourceUsageProvider.setThreadMemoryMeasurementEnabled(true); - ResourceUsageAccountantFactory.ResourceUsageAccountant accountant = - new ResourceUsageAccountantFactory.ResourceUsageAccountant(new PinotConfiguration(configs), - "testWithPerQueryAccountantFactory", InstanceType.SERVER); - - try (MockedStatic tracing = Mockito.mockStatic(Tracing.class, Mockito.CALLS_REAL_METHODS)) { - tracing.when(Tracing::getThreadAccountant).thenReturn(accountant); - ResultTable resultTable = queryRunner("SELECT * FROM a LIMIT 2", false).getResultTable(); - Assert.assertEquals(resultTable.getRows().size(), 2); - - Map resources = accountant.getQueryResources(); - Assert.assertEquals(resources.size(), 1); - Assert.assertEquals(resources.entrySet().iterator().next().getValue().getAllocatedBytes(), 0); - } - } - public static class InterruptingAccountant extends PerQueryCPUMemAccountantFactory.PerQueryCPUMemResourceUsageAccountant { diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceUsageAccountant.java b/pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceUsageAccountant.java index db41723d5bea..b744778e8a93 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceUsageAccountant.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/accounting/ThreadResourceUsageAccountant.java @@ -75,11 +75,6 @@ void setupWorker(int taskId, ThreadExecutionContext.TaskType taskType, */ void sampleUsage(); - /** - * Sample Usage for Multi-stage engine queries - */ - void sampleUsageMSE(); - default boolean throttleQuerySubmission() { return false; } diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/trace/Tracing.java b/pinot-spi/src/main/java/org/apache/pinot/spi/trace/Tracing.java index 6d9993cd4b71..d2466edbe471 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/trace/Tracing.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/trace/Tracing.java @@ -206,10 +206,6 @@ public void clear() { public void sampleUsage() { } - @Override - public void sampleUsageMSE() { - } - @Override public void updateQueryUsageConcurrently(String queryId, long cpuTimeNs, long allocatedBytes, TrackingScope trackingScope) { @@ -295,19 +291,10 @@ public static void sample() { Tracing.getThreadAccountant().sampleUsage(); } - public static void sampleMSE() { - Tracing.getThreadAccountant().sampleUsageMSE(); - } - public static void clear() { Tracing.getThreadAccountant().clear(); } - public static void initializeThreadAccountant(PinotConfiguration config, String instanceId, - InstanceType instanceType) { - createThreadAccountant(config, instanceId, instanceType); - } - public static ThreadResourceUsageAccountant createThreadAccountant(PinotConfiguration config, String instanceId, InstanceType instanceType) { _workloadBudgetManager = new WorkloadBudgetManager(config); diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java index d8152b1812ca..e8c46ff8c8e9 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java @@ -1538,9 +1538,6 @@ public static class Accounting { public static final String CONFIG_OF_QUERY_KILLED_METRIC_ENABLED = "accounting.query.killed.metric.enabled"; public static final boolean DEFAULT_QUERY_KILLED_METRIC_ENABLED = false; - public static final String CONFIG_OF_ENABLE_THREAD_SAMPLING_MSE = "accounting.enable.thread.sampling.mse.debug"; - public static final Boolean DEFAULT_ENABLE_THREAD_SAMPLING_MSE = true; - public static final String CONFIG_OF_CANCEL_CALLBACK_CACHE_MAX_SIZE = "accounting.cancel.callback.cache.max.size"; public static final int DEFAULT_CANCEL_CALLBACK_CACHE_MAX_SIZE = 500; diff --git a/pinot-spi/src/test/java/org/apache/pinot/spi/executor/ThrottleOnCriticalHeapUsageExecutorTest.java b/pinot-spi/src/test/java/org/apache/pinot/spi/executor/ThrottleOnCriticalHeapUsageExecutorTest.java index 141b55129f0b..5005c5436775 100644 --- a/pinot-spi/src/test/java/org/apache/pinot/spi/executor/ThrottleOnCriticalHeapUsageExecutorTest.java +++ b/pinot-spi/src/test/java/org/apache/pinot/spi/executor/ThrottleOnCriticalHeapUsageExecutorTest.java @@ -75,10 +75,6 @@ public ThreadExecutionContext getThreadExecutionContext() { public void sampleUsage() { } - @Override - public void sampleUsageMSE() { - } - @Override public boolean throttleQuerySubmission() { return _numCalls.getAndIncrement() > 1; From 0b0b35c8be3b6896ee53e14d42d59fc7b536a190 Mon Sep 17 00:00:00 2001 From: 9aman <35227405+9aman@users.noreply.github.com> Date: Mon, 4 Aug 2025 21:59:03 +0530 Subject: [PATCH 067/167] Provide the correct comparison time for evaluating whether a segment should be considered in the SegmentStatusChecker (#16496) --- .../pinot/controller/helix/SegmentStatusChecker.java | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/SegmentStatusChecker.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/SegmentStatusChecker.java index 9b16f66288ec..553dada84f87 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/SegmentStatusChecker.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/SegmentStatusChecker.java @@ -287,6 +287,7 @@ private void updateSegmentMetrics(String tableNameWithType, TableConfig tableCon return; } + long evSnapshotTimestamp = System.currentTimeMillis(); ExternalView externalView = _pinotHelixResourceManager.getTableExternalView(tableNameWithType); if (externalView != null) { _controllerMetrics.setValueOfTableGauge(tableNameWithType, ControllerGauge.EXTERNALVIEW_ZNODE_SIZE, @@ -349,9 +350,13 @@ private void updateSegmentMetrics(String tableNameWithType, TableConfig tableCon // the ZK metadata; for pushed segments, we use push time from the ZK metadata. Both of them are the time // when segment is newly created. For committed segments from real-time table, push time doesn't exist, and // creationTimeMs will be Long.MIN_VALUE, which is fine because we want to include them in the check. + // The comparison uses evSnapshotTimestamp instead of System.currentTimeMillis() because for large tables + // with many segments, the status check can take several minutes. A segment updated after + // the EV snapshot was taken but before this individual segment check runs could be incorrectly flagged as + // OFFLINE when using current time. long creationTimeMs = segmentZKMetadata.getStatus() == Status.IN_PROGRESS ? segmentZKMetadata.getCreationTime() : segmentZKMetadata.getPushTime(); - if (creationTimeMs > System.currentTimeMillis() - _waitForPushTimeSeconds * 1000L) { + if (creationTimeMs > evSnapshotTimestamp - _waitForPushTimeSeconds * 1000L) { continue; } From c867974e902505dd30b0d24f85cead8439d2d2e6 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Mon, 4 Aug 2025 13:53:44 -0700 Subject: [PATCH 068/167] Bump com.azure:azure-sdk-bom from 1.2.36 to 1.2.37 (#16501) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 82b9f3514802..2d1db8f0741d 100644 --- a/pom.xml +++ b/pom.xml @@ -179,7 +179,7 @@ 0.4.7 4.2.2 2.32.13 - 1.2.36 + 1.2.37 1.22.0 2.14.0 3.1.12 From 7750e0397bc5fb537eae16aa1ba8640a8b110bec Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Mon, 4 Aug 2025 13:59:22 -0700 Subject: [PATCH 069/167] Bump software.amazon.awssdk:bom from 2.32.13 to 2.32.14 (#16499) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 2d1db8f0741d..b68efd7ccce0 100644 --- a/pom.xml +++ b/pom.xml @@ -178,7 +178,7 @@ 0.15.1 0.4.7 4.2.2 - 2.32.13 + 2.32.14 1.2.37 1.22.0 2.14.0 From 12a3d2938c0c8f8ed5c5eb52a711c7d313446889 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Mon, 4 Aug 2025 13:59:37 -0700 Subject: [PATCH 070/167] Bump commons-cli:commons-cli from 1.9.0 to 1.10.0 (#16502) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index b68efd7ccce0..c359a04a3073 100644 --- a/pom.xml +++ b/pom.xml @@ -212,7 +212,7 @@ 1.11.0 2.20.0 1.19.0 - 1.9.0 + 1.10.0 3.11.1 1.10.0 From 5ceec5d6b1326f2ba2a8611be53728616bdfdb6e Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Mon, 4 Aug 2025 14:06:35 -0700 Subject: [PATCH 071/167] Bump net.minidev:json-smart from 2.5.2 to 2.6.0 (#16500) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index c359a04a3073..66d0f36dca0a 100644 --- a/pom.xml +++ b/pom.xml @@ -154,7 +154,7 @@ 5.27.0 3.4.1 2.9.0 - 2.5.2 + 2.6.0 2.5.0 1.40.0 2.11.1 From 923e61d6fe7143e3426fe208eccfd66d8a36dfb8 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Mon, 4 Aug 2025 14:07:04 -0700 Subject: [PATCH 072/167] Bump commons-net:commons-net from 3.11.1 to 3.12.0 (#16503) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 66d0f36dca0a..380becfd2aa7 100644 --- a/pom.xml +++ b/pom.xml @@ -213,7 +213,7 @@ 2.20.0 1.19.0 1.10.0 - 3.11.1 + 3.12.0 1.10.0 From 4abe567109d9547fc6a29818a0a17a374df007de Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Mon, 4 Aug 2025 14:07:49 -0700 Subject: [PATCH 073/167] Bump com.google.cloud:libraries-bom from 26.64.0 to 26.65.0 (#16504) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 380becfd2aa7..4c9319a2b0eb 100644 --- a/pom.xml +++ b/pom.xml @@ -240,7 +240,7 @@ 3.25.8 1.74.0 - 26.64.0 + 26.65.0 1.1.1 1.8 2.41.0 From 4f46771711d260427a4b8267fc3b9f41492d2f6f Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Mon, 4 Aug 2025 15:10:59 -0700 Subject: [PATCH 074/167] Fixing typo CONFIG_OF_ZOOKEEPER_SERVER in the codebase (#16505) --- .../broker/broker/helix/BaseBrokerStarter.java | 2 +- .../broker/broker/helix/HelixBrokerStarter.java | 4 ++-- .../HelixBrokerStarterHostnamePortTest.java | 10 +++++----- .../broker/broker/HelixBrokerStarterTest.java | 2 +- .../apache/pinot/controller/ControllerConf.java | 6 +++--- .../pinot/integration/tests/ClusterTest.java | 6 +++--- .../tests/ServerStarterIntegrationTest.java | 2 +- .../java/org/apache/pinot/minion/MinionConf.java | 2 +- .../org/apache/pinot/minion/MinionStarter.java | 2 +- .../server/predownload/PredownloadScheduler.java | 2 +- .../server/starter/helix/BaseServerStarter.java | 2 +- .../server/starter/helix/HelixServerStarter.java | 4 ++-- .../apache/pinot/spi/utils/CommonConstants.java | 2 ++ .../tools/admin/command/StartBrokerCommand.java | 2 +- .../tools/admin/command/StartMinionCommand.java | 2 +- .../tools/admin/command/StartServerCommand.java | 2 +- .../pinot/tools/perf/PerfBenchmarkDriver.java | 6 +++--- .../pinot/tools/service/PinotServiceManager.java | 16 ++++++++-------- .../pinot/tools/utils/PinotConfigUtils.java | 6 +++--- 19 files changed, 41 insertions(+), 39 deletions(-) diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java index b127e61a6d54..20646204f86e 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java @@ -161,7 +161,7 @@ public void init(PinotConfiguration brokerConf) throws Exception { _brokerConf = brokerConf; // Remove all white-spaces from the list of zkServers (if any). - _zkServers = brokerConf.getProperty(Helix.CONFIG_OF_ZOOKEEPR_SERVER).replaceAll("\\s+", ""); + _zkServers = brokerConf.getProperty(Helix.CONFIG_OF_ZOOKEEPER_SERVER).replaceAll("\\s+", ""); _clusterName = brokerConf.getProperty(Helix.CONFIG_OF_CLUSTER_NAME); ServiceStartableUtils.applyClusterConfig(_brokerConf, _zkServers, _clusterName, ServiceRole.BROKER); applyCustomConfigs(brokerConf); diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/HelixBrokerStarter.java b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/HelixBrokerStarter.java index fefcd8b38b7a..0de727f44e23 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/HelixBrokerStarter.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/HelixBrokerStarter.java @@ -53,7 +53,7 @@ public HelixBrokerStarter(PinotConfiguration brokerConf, String clusterName, Str private static PinotConfiguration applyBrokerConfigs(PinotConfiguration brokerConf, String clusterName, String zkServers, @Nullable String brokerHost) { brokerConf.setProperty(Helix.CONFIG_OF_CLUSTER_NAME, clusterName); - brokerConf.setProperty(Helix.CONFIG_OF_ZOOKEEPR_SERVER, zkServers); + brokerConf.setProperty(Helix.CONFIG_OF_ZOOKEEPER_SERVER, zkServers); if (brokerHost == null) { brokerConf.clearProperty(Broker.CONFIG_OF_BROKER_HOSTNAME); } else { @@ -75,7 +75,7 @@ public static HelixBrokerStarter getDefault() properties.put(Helix.KEY_OF_BROKER_QUERY_PORT, 5001); properties.put(Broker.CONFIG_OF_BROKER_TIMEOUT_MS, 60 * 1000L); properties.put(Helix.CONFIG_OF_CLUSTER_NAME, "quickstart"); - properties.put(Helix.CONFIG_OF_ZOOKEEPR_SERVER, "localhost:2122"); + properties.put(Helix.CONFIG_OF_ZOOKEEPER_SERVER, "localhost:2122"); HelixBrokerStarter helixBrokerStarter = new HelixBrokerStarter(); helixBrokerStarter.init(new PinotConfiguration(properties)); return helixBrokerStarter; diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/broker/HelixBrokerStarterHostnamePortTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/broker/HelixBrokerStarterHostnamePortTest.java index d2e2d6014081..554368b4d064 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/broker/HelixBrokerStarterHostnamePortTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/broker/HelixBrokerStarterHostnamePortTest.java @@ -33,7 +33,7 @@ import static org.apache.pinot.spi.utils.CommonConstants.Broker.CONFIG_OF_BROKER_ID; import static org.apache.pinot.spi.utils.CommonConstants.Broker.CONFIG_OF_DELAY_SHUTDOWN_TIME_MS; import static org.apache.pinot.spi.utils.CommonConstants.Helix.CONFIG_OF_CLUSTER_NAME; -import static org.apache.pinot.spi.utils.CommonConstants.Helix.CONFIG_OF_ZOOKEEPR_SERVER; +import static org.apache.pinot.spi.utils.CommonConstants.Helix.CONFIG_OF_ZOOKEEPER_SERVER; import static org.apache.pinot.spi.utils.CommonConstants.Helix.Instance.INSTANCE_ID_KEY; import static org.apache.pinot.spi.utils.CommonConstants.Helix.KEY_OF_BROKER_QUERY_PORT; import static org.testng.Assert.assertEquals; @@ -52,7 +52,7 @@ public void setUp() public void testHostnamePortOverride() throws Exception { Map properties = new HashMap<>(); - properties.put(CONFIG_OF_ZOOKEEPR_SERVER, getZkUrl()); + properties.put(CONFIG_OF_ZOOKEEPER_SERVER, getZkUrl()); properties.put(CONFIG_OF_CLUSTER_NAME, getHelixClusterName()); properties.put(INSTANCE_ID_KEY, "Broker_myInstance"); properties.put(CONFIG_OF_BROKER_HOSTNAME, "myHost"); @@ -77,7 +77,7 @@ public void testHostnamePortOverride() public void testInvalidInstanceId() throws Exception { Map properties = new HashMap<>(); - properties.put(CONFIG_OF_ZOOKEEPR_SERVER, getZkUrl()); + properties.put(CONFIG_OF_ZOOKEEPER_SERVER, getZkUrl()); properties.put(CONFIG_OF_CLUSTER_NAME, getHelixClusterName()); properties.put(INSTANCE_ID_KEY, "myInstance"); properties.put(CONFIG_OF_BROKER_HOSTNAME, "myHost"); @@ -91,7 +91,7 @@ public void testInvalidInstanceId() public void testDefaultInstanceId() throws Exception { Map properties = new HashMap<>(); - properties.put(CONFIG_OF_ZOOKEEPR_SERVER, getZkUrl()); + properties.put(CONFIG_OF_ZOOKEEPER_SERVER, getZkUrl()); properties.put(CONFIG_OF_CLUSTER_NAME, getHelixClusterName()); properties.put(CONFIG_OF_BROKER_HOSTNAME, "myHost"); properties.put(KEY_OF_BROKER_QUERY_PORT, 1234); @@ -116,7 +116,7 @@ public void testInstanceIdPrecedence() throws Exception { // Ensures that pinot.broker.instance.id has higher precedence compared to instanceId Map properties = new HashMap<>(); - properties.put(CONFIG_OF_ZOOKEEPR_SERVER, getZkUrl()); + properties.put(CONFIG_OF_ZOOKEEPER_SERVER, getZkUrl()); properties.put(CONFIG_OF_CLUSTER_NAME, getHelixClusterName()); properties.put(CONFIG_OF_BROKER_ID, "Broker_morePrecedence"); properties.put(INSTANCE_ID_KEY, "Broker_lessPrecedence"); diff --git a/pinot-broker/src/test/java/org/apache/pinot/broker/broker/HelixBrokerStarterTest.java b/pinot-broker/src/test/java/org/apache/pinot/broker/broker/HelixBrokerStarterTest.java index 3310a0e8959b..6052d8a663de 100644 --- a/pinot-broker/src/test/java/org/apache/pinot/broker/broker/HelixBrokerStarterTest.java +++ b/pinot-broker/src/test/java/org/apache/pinot/broker/broker/HelixBrokerStarterTest.java @@ -84,7 +84,7 @@ public void setUp() Map properties = new HashMap<>(); properties.put(Helix.KEY_OF_BROKER_QUERY_PORT, 18099); properties.put(Helix.CONFIG_OF_CLUSTER_NAME, getHelixClusterName()); - properties.put(Helix.CONFIG_OF_ZOOKEEPR_SERVER, getZkUrl()); + properties.put(Helix.CONFIG_OF_ZOOKEEPER_SERVER, getZkUrl()); properties.put(Broker.CONFIG_OF_ENABLE_QUERY_LIMIT_OVERRIDE, true); properties.put(Broker.CONFIG_OF_DELAY_SHUTDOWN_TIME_MS, 0); properties.put(Broker.CONFIG_OF_BROKER_DEFAULT_QUERY_LIMIT, 1000); diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/ControllerConf.java b/pinot-controller/src/main/java/org/apache/pinot/controller/ControllerConf.java index 67a48e679f05..c90f26173635 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/ControllerConf.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/ControllerConf.java @@ -485,7 +485,7 @@ public void setMinISSizeForCompression(int minSize) { } public void setZkStr(String zkStr) { - setProperty(CommonConstants.Helix.CONFIG_OF_ZOOKEEPR_SERVER, zkStr); + setProperty(CommonConstants.Helix.CONFIG_OF_ZOOKEEPER_SERVER, zkStr); } public void setDimTableMaxSize(String size) { @@ -565,8 +565,8 @@ public String generateVipUrl() { } public String getZkStr() { - String zkAddress = containsKey(CommonConstants.Helix.CONFIG_OF_ZOOKEEPR_SERVER) ? getProperty( - CommonConstants.Helix.CONFIG_OF_ZOOKEEPR_SERVER) : getProperty(ZK_STR); + String zkAddress = containsKey(CommonConstants.Helix.CONFIG_OF_ZOOKEEPER_SERVER) ? getProperty( + CommonConstants.Helix.CONFIG_OF_ZOOKEEPER_SERVER) : getProperty(ZK_STR); Preconditions.checkState(zkAddress != null, "ZK address is not configured. Please configure it using the config: 'pinot.zk.server'"); return zkAddress; diff --git a/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/ClusterTest.java b/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/ClusterTest.java index 517f3bb8e547..70e131be55d4 100644 --- a/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/ClusterTest.java +++ b/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/ClusterTest.java @@ -173,7 +173,7 @@ protected BaseBrokerStarter createBrokerStarter() { protected PinotConfiguration getBrokerConf(int brokerId) { PinotConfiguration brokerConf = new PinotConfiguration(); - brokerConf.setProperty(Helix.CONFIG_OF_ZOOKEEPR_SERVER, getZkUrl()); + brokerConf.setProperty(Helix.CONFIG_OF_ZOOKEEPER_SERVER, getZkUrl()); brokerConf.setProperty(Helix.CONFIG_OF_CLUSTER_NAME, getHelixClusterName()); brokerConf.setProperty(Broker.CONFIG_OF_BROKER_HOSTNAME, LOCAL_HOST); int brokerPort = NetUtils.findOpenPort(_nextBrokerPort); @@ -240,7 +240,7 @@ protected BaseServerStarter createServerStarter() { protected PinotConfiguration getServerConf(int serverId) { PinotConfiguration serverConf = new PinotConfiguration(); - serverConf.setProperty(Helix.CONFIG_OF_ZOOKEEPR_SERVER, getZkUrl()); + serverConf.setProperty(Helix.CONFIG_OF_ZOOKEEPER_SERVER, getZkUrl()); serverConf.setProperty(Helix.CONFIG_OF_CLUSTER_NAME, getHelixClusterName()); serverConf.setProperty(Helix.KEY_OF_SERVER_NETTY_HOST, LOCAL_HOST); serverConf.setProperty(Server.CONFIG_OF_INSTANCE_DATA_DIR, @@ -315,7 +315,7 @@ protected BaseMinionStarter createMinionStarter() { protected PinotConfiguration getMinionConf() { PinotConfiguration minionConf = new PinotConfiguration(); - minionConf.setProperty(Helix.CONFIG_OF_ZOOKEEPR_SERVER, getZkUrl()); + minionConf.setProperty(Helix.CONFIG_OF_ZOOKEEPER_SERVER, getZkUrl()); minionConf.setProperty(Helix.CONFIG_OF_CLUSTER_NAME, getHelixClusterName()); minionConf.setProperty(Helix.KEY_OF_MINION_HOST, LOCAL_HOST); int minionPort = NetUtils.findOpenPort(_nextMinionPort); diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/ServerStarterIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/ServerStarterIntegrationTest.java index 76ddb1e1e7c4..5aa9cd58f739 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/ServerStarterIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/ServerStarterIntegrationTest.java @@ -56,7 +56,7 @@ private void verifyInstanceConfig(PinotConfiguration serverConf, String expected int expectedPort) throws Exception { serverConf.setProperty(CONFIG_OF_CLUSTER_NAME, getHelixClusterName()); - serverConf.setProperty(CONFIG_OF_ZOOKEEPR_SERVER, getZkUrl()); + serverConf.setProperty(CONFIG_OF_ZOOKEEPER_SERVER, getZkUrl()); HelixServerStarter helixServerStarter = new HelixServerStarter(); helixServerStarter.init(serverConf); helixServerStarter.start(); diff --git a/pinot-minion/src/main/java/org/apache/pinot/minion/MinionConf.java b/pinot-minion/src/main/java/org/apache/pinot/minion/MinionConf.java index fde239f654ec..019f20e09457 100644 --- a/pinot-minion/src/main/java/org/apache/pinot/minion/MinionConf.java +++ b/pinot-minion/src/main/java/org/apache/pinot/minion/MinionConf.java @@ -51,7 +51,7 @@ public String getHelixClusterName() { } public String getZkAddress() { - return getProperty(CommonConstants.Helix.CONFIG_OF_ZOOKEEPR_SERVER); + return getProperty(CommonConstants.Helix.CONFIG_OF_ZOOKEEPER_SERVER); } public String getHostName() diff --git a/pinot-minion/src/main/java/org/apache/pinot/minion/MinionStarter.java b/pinot-minion/src/main/java/org/apache/pinot/minion/MinionStarter.java index cc0f2a25d6e5..5e050fb6b05e 100644 --- a/pinot-minion/src/main/java/org/apache/pinot/minion/MinionStarter.java +++ b/pinot-minion/src/main/java/org/apache/pinot/minion/MinionStarter.java @@ -41,7 +41,7 @@ public MinionStarter(String clusterName, String zkServers, PinotConfiguration mi private static PinotConfiguration applyMinionConfigs(PinotConfiguration minionConfig, String clusterName, String zkServers) { minionConfig.setProperty(CommonConstants.Helix.CONFIG_OF_CLUSTER_NAME, clusterName); - minionConfig.setProperty(CommonConstants.Helix.CONFIG_OF_ZOOKEEPR_SERVER, zkServers); + minionConfig.setProperty(CommonConstants.Helix.CONFIG_OF_ZOOKEEPER_SERVER, zkServers); return minionConfig; } diff --git a/pinot-server/src/main/java/org/apache/pinot/server/predownload/PredownloadScheduler.java b/pinot-server/src/main/java/org/apache/pinot/server/predownload/PredownloadScheduler.java index 5638c14188d8..835ef4016216 100644 --- a/pinot-server/src/main/java/org/apache/pinot/server/predownload/PredownloadScheduler.java +++ b/pinot-server/src/main/java/org/apache/pinot/server/predownload/PredownloadScheduler.java @@ -84,7 +84,7 @@ public PredownloadScheduler(PropertiesConfiguration properties) throws Exception { _properties = properties; _clusterName = properties.getString(CommonConstants.Helix.CONFIG_OF_CLUSTER_NAME); - _zkAddress = properties.getString(CommonConstants.Helix.CONFIG_OF_ZOOKEEPR_SERVER); + _zkAddress = properties.getString(CommonConstants.Helix.CONFIG_OF_ZOOKEEPER_SERVER); _instanceId = properties.getString(CommonConstants.Server.CONFIG_OF_INSTANCE_ID); _pinotConfig = new PinotConfiguration(properties); _instanceDataManagerConfig = diff --git a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/BaseServerStarter.java b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/BaseServerStarter.java index c3f6f0e611ad..764a2c46a48f 100644 --- a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/BaseServerStarter.java +++ b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/BaseServerStarter.java @@ -179,7 +179,7 @@ public void init(PinotConfiguration serverConf) throws Exception { // Make a clone so that changes to the config won't propagate to the caller _serverConf = serverConf.clone(); - _zkAddress = _serverConf.getProperty(CommonConstants.Helix.CONFIG_OF_ZOOKEEPR_SERVER); + _zkAddress = _serverConf.getProperty(CommonConstants.Helix.CONFIG_OF_ZOOKEEPER_SERVER); _helixClusterName = _serverConf.getProperty(CommonConstants.Helix.CONFIG_OF_CLUSTER_NAME); ServiceStartableUtils.applyClusterConfig(_serverConf, _zkAddress, _helixClusterName, ServiceRole.SERVER); applyCustomConfigs(_serverConf); diff --git a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixServerStarter.java b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixServerStarter.java index 7321a0a70988..31996dddd648 100644 --- a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixServerStarter.java +++ b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/HelixServerStarter.java @@ -59,7 +59,7 @@ public HelixServerStarter(String helixClusterName, String zkAddress, PinotConfig private static PinotConfiguration applyServerConfig(PinotConfiguration serverConf, String helixClusterName, String zkAddress) { serverConf.setProperty(Helix.CONFIG_OF_CLUSTER_NAME, helixClusterName); - serverConf.setProperty(Helix.CONFIG_OF_ZOOKEEPR_SERVER, zkAddress); + serverConf.setProperty(Helix.CONFIG_OF_ZOOKEEPER_SERVER, zkAddress); return serverConf; } @@ -77,7 +77,7 @@ public static HelixServerStarter startDefault() Map properties = new HashMap<>(); int port = 8003; properties.put(Helix.CONFIG_OF_CLUSTER_NAME, "quickstart"); - properties.put(Helix.CONFIG_OF_ZOOKEEPR_SERVER, "localhost:2191"); + properties.put(Helix.CONFIG_OF_ZOOKEEPER_SERVER, "localhost:2191"); properties.put(Helix.KEY_OF_SERVER_NETTY_PORT, port); properties.put(Server.CONFIG_OF_INSTANCE_DATA_DIR, "/tmp/PinotServer/test" + port + "/index"); properties.put(Server.CONFIG_OF_INSTANCE_SEGMENT_TAR_DIR, "/tmp/PinotServer/test" + port + "/segmentTar"); diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java index e8c46ff8c8e9..894c73f8dd91 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java @@ -241,6 +241,8 @@ public static class Instance { public static final String DEFAULT_FLAPPING_TIME_WINDOW_MS = "1"; public static final String PINOT_SERVICE_ROLE = "pinot.service.role"; public static final String CONFIG_OF_CLUSTER_NAME = "pinot.cluster.name"; + public static final String CONFIG_OF_ZOOKEEPER_SERVER = "pinot.zk.server"; + @Deprecated(since = "1.5.0", forRemoval = true) public static final String CONFIG_OF_ZOOKEEPR_SERVER = "pinot.zk.server"; public static final String CONFIG_OF_PINOT_CONTROLLER_STARTABLE_CLASS = "pinot.controller.startable.class"; diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/StartBrokerCommand.java b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/StartBrokerCommand.java index 9ec09f9c78af..95d542680f3b 100644 --- a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/StartBrokerCommand.java +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/StartBrokerCommand.java @@ -176,7 +176,7 @@ protected Map getBrokerConf() properties.putAll(PinotConfigUtils.readConfigFromFile(_configFileName)); // Override the zkAddress and clusterName to ensure ServiceManager is connecting to the right Zookeeper and // Cluster. - _zkAddress = MapUtils.getString(properties, CommonConstants.Helix.CONFIG_OF_ZOOKEEPR_SERVER, _zkAddress); + _zkAddress = MapUtils.getString(properties, CommonConstants.Helix.CONFIG_OF_ZOOKEEPER_SERVER, _zkAddress); _clusterName = MapUtils.getString(properties, CommonConstants.Helix.CONFIG_OF_CLUSTER_NAME, _clusterName); } else { properties.putAll(PinotConfigUtils.generateBrokerConf(_clusterName, _zkAddress, _brokerHost, _brokerPort, diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/StartMinionCommand.java b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/StartMinionCommand.java index 934e1de10c2f..7c8613b63354 100644 --- a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/StartMinionCommand.java +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/StartMinionCommand.java @@ -136,7 +136,7 @@ protected Map getMinionConf() properties.putAll(PinotConfigUtils.readConfigFromFile(_configFileName)); // Override the zkAddress and clusterName to ensure ServiceManager is connecting to the right Zookeeper and // Cluster. - _zkAddress = MapUtils.getString(properties, CommonConstants.Helix.CONFIG_OF_ZOOKEEPR_SERVER, _zkAddress); + _zkAddress = MapUtils.getString(properties, CommonConstants.Helix.CONFIG_OF_ZOOKEEPER_SERVER, _zkAddress); _clusterName = MapUtils.getString(properties, CommonConstants.Helix.CONFIG_OF_CLUSTER_NAME, _clusterName); } else { properties.putAll(PinotConfigUtils.generateMinionConf(_clusterName, _zkAddress, _minionHost, _minionPort)); diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/StartServerCommand.java b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/StartServerCommand.java index e091f88de0d4..3e76ab11d7f6 100644 --- a/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/StartServerCommand.java +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/admin/command/StartServerCommand.java @@ -246,7 +246,7 @@ protected Map getServerConf() properties.putAll(PinotConfigUtils.readConfigFromFile(_configFileName)); // Override the zkAddress and clusterName to ensure ServiceManager is connecting to the right Zookeeper and // Cluster. - _zkAddress = MapUtils.getString(properties, CommonConstants.Helix.CONFIG_OF_ZOOKEEPR_SERVER, _zkAddress); + _zkAddress = MapUtils.getString(properties, CommonConstants.Helix.CONFIG_OF_ZOOKEEPER_SERVER, _zkAddress); _clusterName = MapUtils.getString(properties, CommonConstants.Helix.CONFIG_OF_CLUSTER_NAME, _clusterName); } else { properties.putAll(PinotConfigUtils diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/perf/PerfBenchmarkDriver.java b/pinot-tools/src/main/java/org/apache/pinot/tools/perf/PerfBenchmarkDriver.java index 526b500854c0..c6f93bda71a3 100644 --- a/pinot-tools/src/main/java/org/apache/pinot/tools/perf/PerfBenchmarkDriver.java +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/perf/PerfBenchmarkDriver.java @@ -212,7 +212,7 @@ private void startController() private Map getControllerProperties() { Map properties = new HashMap<>(); properties.put(CommonConstants.Helix.CONFIG_OF_CLUSTER_NAME, _clusterName); - properties.put(CommonConstants.Helix.CONFIG_OF_ZOOKEEPR_SERVER, _zkAddress); + properties.put(CommonConstants.Helix.CONFIG_OF_ZOOKEEPER_SERVER, _zkAddress); properties.put(ControllerConf.CONTROLLER_HOST, _controllerHost); properties.put(ControllerConf.CONTROLLER_PORT, String.valueOf(_controllerPort)); properties.put(ControllerConf.DATA_DIR, _controllerDataDir); @@ -247,7 +247,7 @@ private void startBroker() properties.put(CommonConstants.Broker.CONFIG_OF_BROKER_ID, brokerInstanceName); properties.put(CommonConstants.Broker.CONFIG_OF_BROKER_TIMEOUT_MS, BROKER_TIMEOUT_MS); properties.put(CommonConstants.Helix.CONFIG_OF_CLUSTER_NAME, _clusterName); - properties.put(CommonConstants.Helix.CONFIG_OF_ZOOKEEPR_SERVER, _zkAddress); + properties.put(CommonConstants.Helix.CONFIG_OF_ZOOKEEPER_SERVER, _zkAddress); LOGGER.info("Starting broker instance: {}", brokerInstanceName); @@ -269,7 +269,7 @@ private void startServer() properties.put(CommonConstants.Helix.KEY_OF_SERVER_NETTY_HOST, "localhost"); properties.put(CommonConstants.Server.CONFIG_OF_INSTANCE_ID, _serverInstanceName); properties.put(CommonConstants.Helix.CONFIG_OF_CLUSTER_NAME, _clusterName); - properties.put(CommonConstants.Helix.CONFIG_OF_ZOOKEEPR_SERVER, _zkAddress); + properties.put(CommonConstants.Helix.CONFIG_OF_ZOOKEEPER_SERVER, _zkAddress); if (_segmentFormatVersion != null) { properties.put(CommonConstants.Server.CONFIG_OF_SEGMENT_FORMAT_VERSION, _segmentFormatVersion); } diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/service/PinotServiceManager.java b/pinot-tools/src/main/java/org/apache/pinot/tools/service/PinotServiceManager.java index 5ecc7e007edf..f4ff1ef2f627 100644 --- a/pinot-tools/src/main/java/org/apache/pinot/tools/service/PinotServiceManager.java +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/service/PinotServiceManager.java @@ -110,8 +110,8 @@ public String startController(String controllerStarterClassName, PinotConfigurat if (!controllerConf.containsKey(CommonConstants.Helix.CONFIG_OF_CLUSTER_NAME)) { controllerConf.setProperty(CommonConstants.Helix.CONFIG_OF_CLUSTER_NAME, _clusterName); } - if (!controllerConf.containsKey(CommonConstants.Helix.CONFIG_OF_ZOOKEEPR_SERVER)) { - controllerConf.setProperty(CommonConstants.Helix.CONFIG_OF_ZOOKEEPR_SERVER, _zkAddress); + if (!controllerConf.containsKey(CommonConstants.Helix.CONFIG_OF_ZOOKEEPER_SERVER)) { + controllerConf.setProperty(CommonConstants.Helix.CONFIG_OF_ZOOKEEPER_SERVER, _zkAddress); } ServiceStartable controllerStarter = getServiceStartable(controllerStarterClassName); controllerStarter.init(controllerConf); @@ -128,8 +128,8 @@ public String startBroker(String brokerStarterClassName, PinotConfiguration brok if (!brokerConf.containsKey(CommonConstants.Helix.CONFIG_OF_CLUSTER_NAME)) { brokerConf.setProperty(CommonConstants.Helix.CONFIG_OF_CLUSTER_NAME, _clusterName); } - if (!brokerConf.containsKey(CommonConstants.Helix.CONFIG_OF_ZOOKEEPR_SERVER)) { - brokerConf.setProperty(CommonConstants.Helix.CONFIG_OF_ZOOKEEPR_SERVER, _zkAddress); + if (!brokerConf.containsKey(CommonConstants.Helix.CONFIG_OF_ZOOKEEPER_SERVER)) { + brokerConf.setProperty(CommonConstants.Helix.CONFIG_OF_ZOOKEEPER_SERVER, _zkAddress); } ServiceStartable brokerStarter; try { @@ -159,8 +159,8 @@ public String startServer(String serverStarterClassName, PinotConfiguration serv serverConf.setProperty(CommonConstants.Helix.CONFIG_OF_CLUSTER_NAME, _clusterName); } - if (!serverConf.containsKey(CommonConstants.Helix.CONFIG_OF_ZOOKEEPR_SERVER)) { - serverConf.setProperty(CommonConstants.Helix.CONFIG_OF_ZOOKEEPR_SERVER, _zkAddress); + if (!serverConf.containsKey(CommonConstants.Helix.CONFIG_OF_ZOOKEEPER_SERVER)) { + serverConf.setProperty(CommonConstants.Helix.CONFIG_OF_ZOOKEEPER_SERVER, _zkAddress); } ServiceStartable serverStarter = getServiceStartable(serverStarterClassName); serverStarter.init(serverConf); @@ -178,8 +178,8 @@ public String startMinion(String minionStarterClassName, PinotConfiguration mini if (!minionConf.containsKey(CommonConstants.Helix.CONFIG_OF_CLUSTER_NAME)) { minionConf.setProperty(CommonConstants.Helix.CONFIG_OF_CLUSTER_NAME, _clusterName); } - if (!minionConf.containsKey(CommonConstants.Helix.CONFIG_OF_ZOOKEEPR_SERVER)) { - minionConf.setProperty(CommonConstants.Helix.CONFIG_OF_ZOOKEEPR_SERVER, _zkAddress); + if (!minionConf.containsKey(CommonConstants.Helix.CONFIG_OF_ZOOKEEPER_SERVER)) { + minionConf.setProperty(CommonConstants.Helix.CONFIG_OF_ZOOKEEPER_SERVER, _zkAddress); } ServiceStartable minionStarter = getServiceStartable(minionStarterClassName); minionStarter.init(minionConf); diff --git a/pinot-tools/src/main/java/org/apache/pinot/tools/utils/PinotConfigUtils.java b/pinot-tools/src/main/java/org/apache/pinot/tools/utils/PinotConfigUtils.java index 0a9e1cd31eae..d2571faafd7d 100644 --- a/pinot-tools/src/main/java/org/apache/pinot/tools/utils/PinotConfigUtils.java +++ b/pinot-tools/src/main/java/org/apache/pinot/tools/utils/PinotConfigUtils.java @@ -156,7 +156,7 @@ public static Map generateBrokerConf(String clusterName, String throws SocketException, UnknownHostException { Map properties = new HashMap<>(); properties.put(CommonConstants.Helix.CONFIG_OF_CLUSTER_NAME, clusterName); - properties.put(CommonConstants.Helix.CONFIG_OF_ZOOKEEPR_SERVER, zkAddress); + properties.put(CommonConstants.Helix.CONFIG_OF_ZOOKEEPER_SERVER, zkAddress); properties.put(CommonConstants.Broker.CONFIG_OF_BROKER_HOSTNAME, !StringUtils.isEmpty(brokerHost) ? brokerHost : NetUtils.getHostAddress()); properties.put(CommonConstants.Helix.KEY_OF_BROKER_QUERY_PORT, brokerPort != 0 ? brokerPort : getAvailablePort()); @@ -186,7 +186,7 @@ public static Map generateServerConf(String clusterName, String } Map properties = new HashMap<>(); properties.put(CommonConstants.Helix.CONFIG_OF_CLUSTER_NAME, clusterName); - properties.put(CommonConstants.Helix.CONFIG_OF_ZOOKEEPR_SERVER, zkAddress); + properties.put(CommonConstants.Helix.CONFIG_OF_ZOOKEEPER_SERVER, zkAddress); properties.put(CommonConstants.Helix.KEY_OF_SERVER_NETTY_HOST, serverHost); properties.put(CommonConstants.Helix.KEY_OF_SERVER_NETTY_PORT, serverPort); properties.put(CommonConstants.MultiStageQueryRunner.KEY_OF_QUERY_SERVER_PORT, serverMultiStageServerPort != 0 @@ -209,7 +209,7 @@ public static Map generateMinionConf(String clusterName, String } Map properties = new HashMap<>(); properties.put(CommonConstants.Helix.CONFIG_OF_CLUSTER_NAME, clusterName); - properties.put(CommonConstants.Helix.CONFIG_OF_ZOOKEEPR_SERVER, zkAddress); + properties.put(CommonConstants.Helix.CONFIG_OF_ZOOKEEPER_SERVER, zkAddress); properties.put(CommonConstants.Helix.KEY_OF_MINION_HOST, minionHost); properties.put(CommonConstants.Helix.KEY_OF_MINION_PORT, minionPort != 0 ? minionPort : getAvailablePort()); From e061d047b4ec3d13a565ffa61edf53df60196abb Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Tue, 5 Aug 2025 06:07:46 -0700 Subject: [PATCH 075/167] Support time zone in controller ui and add start/end time of segment in the summary page (#16489) --- pinot-controller/src/main/resources/.eslintrc | 8 +- .../src/main/resources/app/App.tsx | 55 +- .../src/main/resources/app/app_state.ts | 5 +- .../app/components/ConsumingSegmentsTable.tsx | 2 +- .../resources/app/components/CustomButton.tsx | 2 +- .../resources/app/components/CustomDialog.tsx | 2 +- .../app/components/CustomMultiSelect.tsx | 2 +- .../app/components/CustomNotification.tsx | 4 +- .../main/resources/app/components/Header.tsx | 34 +- .../app/components/Homepage/ClusterConfig.tsx | 2 +- .../app/components/Homepage/InstanceTable.tsx | 2 +- .../Operations/AddDeleteComponent.tsx | 2 +- .../Operations/AddIndexingComponent.tsx | 2 +- .../Operations/AddIngestionComponent.tsx | 2 +- .../Homepage/Operations/AddOfflineTableOp.tsx | 2 +- .../Operations/AddOfflineTenantComponent.tsx | 2 +- .../Operations/AddPartionComponent.tsx | 2 +- .../Homepage/Operations/AddQueryComponent.tsx | 2 +- .../AddRealTimeIngestionComponent.tsx | 2 +- .../AddRealTimePartionComponent.tsx | 2 +- .../Operations/AddRealtimeTableOp.tsx | 2 +- .../Homepage/Operations/AddSchemaOp.tsx | 16 +- .../Operations/AddStorageComponent.tsx | 2 +- .../Homepage/Operations/AddTableComponent.tsx | 4 +- .../Operations/AddTenantComponent.tsx | 2 +- .../Homepage/Operations/EditConfigOp.tsx | 2 +- .../Homepage/Operations/EditTagsOp.tsx | 2 +- .../Operations/MultiIndexingComponent.tsx | 2 +- .../Operations/MultiMetricComponent.tsx | 2 +- .../Operations/MultipleSelectComponent.tsx | 2 +- .../RebalanceServer/RebalanceResponse.tsx | 2 +- .../RebalanceServerConfigurationOption.tsx | 2 +- ...balanceServerConfigurationOptionDouble.tsx | 2 +- ...alanceServerConfigurationOptionInteger.tsx | 2 +- ...ebalanceServerConfigurationOptionLabel.tsx | 2 +- ...balanceServerConfigurationOptionSelect.tsx | 2 +- ...eServerConfigurationOptionToggleSwitch.tsx | 2 +- .../RebalanceServerDialogHeader.tsx | 2 +- .../RebalanceServer/RebalanceServerOptions.ts | 2 +- .../RebalanceServerPreChecksResponse.tsx | 2 +- ...ebalanceServerRebalanceSummaryResponse.tsx | 2 +- .../RebalanceServerResponseCard.tsx | 2 +- .../RebalanceServerSectionResponse.tsx | 2 +- .../RebalanceServerSection.tsx | 2 +- .../Operations/RebalanceServerStatusOp.tsx | 7 +- .../Operations/RebalanceServerTableOp.tsx | 4 +- .../Operations/RebalanceServerTenantOp.tsx | 4 +- .../Homepage/Operations/ReloadStatusOp.tsx | 17 +- .../Homepage/Operations/SchemaComponent.tsx | 2 +- .../Operations/SchemaNameComponent.tsx | 2 +- .../components/Homepage/useTaskTypesTable.tsx | 2 +- .../main/resources/app/components/Layout.tsx | 3 +- .../Notification/NotificationContext.tsx | 2 +- .../NotificationContextProvider.tsx | 4 +- .../app/components/Query/MetricStatsTable.tsx | 1 - .../Query/VisualizeQueryStageStats.tsx | 2 +- .../resources/app/components/SearchBar.tsx | 2 +- .../app/components/SimpleAccordion.tsx | 2 +- .../resources/app/components/StatusFilter.tsx | 2 +- .../resources/app/components/TabPanel.tsx | 2 +- .../resources/app/components/TableToolbar.tsx | 2 +- .../app/components/TimezoneSelector.tsx | 169 +++++ .../resources/app/components/User/AddUser.tsx | 2 +- .../app/components/User/UpdateUser.tsx | 2 +- .../components/Zookeeper/TreeDirectory.tsx | 2 +- .../app/components/auth/AuthProvider.tsx | 4 +- .../app/components/useMinionMetaData.tsx | 2 +- .../app/components/useScheduleAdhocModal.tsx | 2 +- .../app/components/useTaskListing.tsx | 6 +- .../app/contexts/TimezoneContext.tsx | 66 ++ .../src/main/resources/app/index.tsx | 2 +- .../src/main/resources/app/pages/HomePage.tsx | 8 +- .../app/pages/InstanceListingPage.tsx | 12 +- .../main/resources/app/pages/LoginPage.tsx | 2 +- .../resources/app/pages/MinionTaskManager.tsx | 2 +- .../src/main/resources/app/pages/Query.tsx | 16 +- .../resources/app/pages/SegmentDetails.tsx | 115 +++- .../resources/app/pages/SubTaskDetail.tsx | 33 +- .../main/resources/app/pages/TaskDetail.tsx | 31 +- .../main/resources/app/pages/TaskQueue.tsx | 2 +- .../resources/app/pages/TaskQueueTable.tsx | 13 +- .../resources/app/pages/TenantDetails.tsx | 2 +- .../resources/app/pages/ZookeeperPage.tsx | 2 +- .../src/main/resources/app/requests/index.ts | 10 +- .../src/main/resources/app/router.tsx | 2 +- .../resources/app/utils/PinotMethodUtils.ts | 31 +- .../main/resources/app/utils/TimezoneUtils.ts | 610 ++++++++++++++++++ .../src/main/resources/app/utils/Utils.tsx | 5 +- .../src/main/resources/package-lock.json | 21 + .../src/main/resources/package.json | 3 +- 90 files changed, 1214 insertions(+), 223 deletions(-) create mode 100644 pinot-controller/src/main/resources/app/components/TimezoneSelector.tsx create mode 100644 pinot-controller/src/main/resources/app/contexts/TimezoneContext.tsx create mode 100644 pinot-controller/src/main/resources/app/utils/TimezoneUtils.ts diff --git a/pinot-controller/src/main/resources/.eslintrc b/pinot-controller/src/main/resources/.eslintrc index 2eee0b469860..d505e2023d1c 100644 --- a/pinot-controller/src/main/resources/.eslintrc +++ b/pinot-controller/src/main/resources/.eslintrc @@ -1,6 +1,6 @@ { "plugins": ["prettier", "@typescript-eslint"], - "extends": ["airbnb-typescript", "react-app", "prettier"], + "extends": ["react-app", "prettier"], "parser": "@typescript-eslint/parser", "parserOptions": { "project": "./tsconfig.json" @@ -61,6 +61,10 @@ "react/no-unescaped-entities": "off", "react/jsx-one-expression-per-line": "off", "react/jsx-wrap-multilines": "off", - "react/destructuring-assignment": "off" + "react/destructuring-assignment": "off", + "import/no-cycle": "off", + "no-trailing-spaces": "error", + "no-multiple-empty-lines": ["error", { "max": 2, "maxEOF": 1 }], + "eol-last": ["error", "always"] } } diff --git a/pinot-controller/src/main/resources/app/App.tsx b/pinot-controller/src/main/resources/app/App.tsx index b22ed4143d38..2682c506de36 100644 --- a/pinot-controller/src/main/resources/app/App.tsx +++ b/pinot-controller/src/main/resources/app/App.tsx @@ -26,6 +26,7 @@ import app_state from './app_state'; import { useAuthProvider } from './components/auth/AuthProvider'; import { AppLoadingIndicator } from './components/AppLoadingIndicator'; import { AuthWorkflow } from 'Models'; +import { TimezoneProvider } from './contexts/TimezoneContext'; export const App = () => { const [clusterName, setClusterName] = useState(''); @@ -137,31 +138,33 @@ export const App = () => { } return ( - - {getRouterData().map(({ path, Component }, key) => ( - { - if (path === '/login') { - return loginRender(Component, props); - } else if (isAuthenticated) { - // default render - return componentRender(Component, props, role); - } else { - return ; - } - }} - /> - ))} - - - - + + + {getRouterData().map(({ path, Component }, key) => ( + { + if (path === '/login') { + return loginRender(Component, props); + } else if (isAuthenticated) { + // default render + return componentRender(Component, props, role); + } else { + return ; + } + }} + /> + ))} + + + + + ); }; diff --git a/pinot-controller/src/main/resources/app/app_state.ts b/pinot-controller/src/main/resources/app/app_state.ts index 8b0fa7c3d1d4..27bcdc314f4e 100644 --- a/pinot-controller/src/main/resources/app/app_state.ts +++ b/pinot-controller/src/main/resources/app/app_state.ts @@ -18,6 +18,8 @@ */ import { AuthWorkflow } from "Models"; +import { DEFAULT_TIMEZONE_FALLBACK } from './utils/TimezoneUtils'; + class app_state { queryConsoleOnlyView: boolean; authWorkflow: AuthWorkflow; @@ -26,6 +28,7 @@ class app_state { hideQueryConsoleTab: boolean; username: string; role: string; + timezone: string = DEFAULT_TIMEZONE_FALLBACK; } -export default new app_state(); \ No newline at end of file +export default new app_state(); diff --git a/pinot-controller/src/main/resources/app/components/ConsumingSegmentsTable.tsx b/pinot-controller/src/main/resources/app/components/ConsumingSegmentsTable.tsx index 4b940c17310e..a2b1e44b9687 100644 --- a/pinot-controller/src/main/resources/app/components/ConsumingSegmentsTable.tsx +++ b/pinot-controller/src/main/resources/app/components/ConsumingSegmentsTable.tsx @@ -107,4 +107,4 @@ const ConsumingSegmentsTable: React.FC = ({ info }) => { ); }; -export default ConsumingSegmentsTable; \ No newline at end of file +export default ConsumingSegmentsTable; diff --git a/pinot-controller/src/main/resources/app/components/CustomButton.tsx b/pinot-controller/src/main/resources/app/components/CustomButton.tsx index db6261815008..f212261b0d95 100644 --- a/pinot-controller/src/main/resources/app/components/CustomButton.tsx +++ b/pinot-controller/src/main/resources/app/components/CustomButton.tsx @@ -58,4 +58,4 @@ export default function CustomButton({ ); -} \ No newline at end of file +} diff --git a/pinot-controller/src/main/resources/app/components/CustomDialog.tsx b/pinot-controller/src/main/resources/app/components/CustomDialog.tsx index ca3b659c4040..1896dbbe0322 100644 --- a/pinot-controller/src/main/resources/app/components/CustomDialog.tsx +++ b/pinot-controller/src/main/resources/app/components/CustomDialog.tsx @@ -124,4 +124,4 @@ export default function CustomDialog({ ); -} \ No newline at end of file +} diff --git a/pinot-controller/src/main/resources/app/components/CustomMultiSelect.tsx b/pinot-controller/src/main/resources/app/components/CustomMultiSelect.tsx index 22e3a3a9d756..1ee06ab37323 100644 --- a/pinot-controller/src/main/resources/app/components/CustomMultiSelect.tsx +++ b/pinot-controller/src/main/resources/app/components/CustomMultiSelect.tsx @@ -119,4 +119,4 @@ const CustomMultiSelect = forwardRef(({ /> ); }); -export default CustomMultiSelect; \ No newline at end of file +export default CustomMultiSelect; diff --git a/pinot-controller/src/main/resources/app/components/CustomNotification.tsx b/pinot-controller/src/main/resources/app/components/CustomNotification.tsx index 50c8775fd49d..86f27a195130 100644 --- a/pinot-controller/src/main/resources/app/components/CustomNotification.tsx +++ b/pinot-controller/src/main/resources/app/components/CustomNotification.tsx @@ -29,7 +29,7 @@ const Alert = (props) => { const CustomNotification = () => { return ( - {context => + {context => { ); }; -export default CustomNotification; \ No newline at end of file +export default CustomNotification; diff --git a/pinot-controller/src/main/resources/app/components/Header.tsx b/pinot-controller/src/main/resources/app/components/Header.tsx index 060527d10df0..a5c46a714731 100644 --- a/pinot-controller/src/main/resources/app/components/Header.tsx +++ b/pinot-controller/src/main/resources/app/components/Header.tsx @@ -23,6 +23,7 @@ import { AppBar, Box, makeStyles, Paper } from '@material-ui/core'; import MenuIcon from '@material-ui/icons/Menu'; import Logo from './SvgIcons/Logo'; import BreadcrumbsComponent from './Breadcrumbs'; +import TimezoneSelector from './TimezoneSelector'; type Props = { highlightSidebarLink: (id: number) => void; @@ -53,6 +54,32 @@ const useStyles = makeStyles((theme) => ({ letterSpacing: 1, fontWeight: 500 } + }, + timezoneContainer: { + display: 'flex', + alignItems: 'center', + marginRight: theme.spacing(2), + '& .MuiFormControl-root': { + margin: 0, + }, + '& .MuiInputLabel-root': { + color: 'rgba(255, 255, 255, 0.7)', + }, + '& .MuiSelect-select': { + color: '#fff', + }, + '& .MuiSelect-icon': { + color: 'rgba(255, 255, 255, 0.7)', + }, + '& .MuiOutlinedInput-notchedOutline': { + borderColor: 'rgba(255, 255, 255, 0.3)', + }, + '& .MuiOutlinedInput-root:hover .MuiOutlinedInput-notchedOutline': { + borderColor: 'rgba(255, 255, 255, 0.5)', + }, + '& .MuiOutlinedInput-root.Mui-focused .MuiOutlinedInput-notchedOutline': { + borderColor: 'rgba(255, 255, 255, 0.7)', + }, } })); @@ -70,6 +97,11 @@ const Header = ({ highlightSidebarLink, showHideSideBarHandler, openSidebar, clu + + + + +

    Cluster Name

    @@ -81,4 +113,4 @@ const Header = ({ highlightSidebarLink, showHideSideBarHandler, openSidebar, clu ) }; -export default Header; \ No newline at end of file +export default Header; diff --git a/pinot-controller/src/main/resources/app/components/Homepage/ClusterConfig.tsx b/pinot-controller/src/main/resources/app/components/Homepage/ClusterConfig.tsx index 1b5217c713ec..51ef3f44aa6b 100644 --- a/pinot-controller/src/main/resources/app/components/Homepage/ClusterConfig.tsx +++ b/pinot-controller/src/main/resources/app/components/Homepage/ClusterConfig.tsx @@ -51,4 +51,4 @@ const ClusterConfig = () => { ); }; -export default ClusterConfig; \ No newline at end of file +export default ClusterConfig; diff --git a/pinot-controller/src/main/resources/app/components/Homepage/InstanceTable.tsx b/pinot-controller/src/main/resources/app/components/Homepage/InstanceTable.tsx index 24b45dbeed4a..c04a227505a2 100644 --- a/pinot-controller/src/main/resources/app/components/Homepage/InstanceTable.tsx +++ b/pinot-controller/src/main/resources/app/components/Homepage/InstanceTable.tsx @@ -69,4 +69,4 @@ const InstanceTable = ({ name, instances, clusterName }: Props) => { ); }; -export default InstanceTable; \ No newline at end of file +export default InstanceTable; diff --git a/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddDeleteComponent.tsx b/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddDeleteComponent.tsx index b29833646074..d1c492d077eb 100644 --- a/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddDeleteComponent.tsx +++ b/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddDeleteComponent.tsx @@ -174,4 +174,4 @@ export default function AddDeleteComponent({ ); -} \ No newline at end of file +} diff --git a/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddIndexingComponent.tsx b/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddIndexingComponent.tsx index 78fd5ee37a54..8839c1a9d08c 100644 --- a/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddIndexingComponent.tsx +++ b/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddIndexingComponent.tsx @@ -159,4 +159,4 @@ export default function AddIndexingComponent({ ); -} \ No newline at end of file +} diff --git a/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddIngestionComponent.tsx b/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddIngestionComponent.tsx index 1724b7fd32fd..0d24a0740d91 100644 --- a/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddIngestionComponent.tsx +++ b/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddIngestionComponent.tsx @@ -170,4 +170,4 @@ export default function AddIngestionComponent({ ); -} \ No newline at end of file +} diff --git a/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddOfflineTableOp.tsx b/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddOfflineTableOp.tsx index 3d27f52799c2..266584036cf0 100644 --- a/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddOfflineTableOp.tsx +++ b/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddOfflineTableOp.tsx @@ -449,4 +449,4 @@ const checkFields = (tableObj,fields) => { ); -} \ No newline at end of file +} diff --git a/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddOfflineTenantComponent.tsx b/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddOfflineTenantComponent.tsx index 4b1984d82202..0247b9e69130 100644 --- a/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddOfflineTenantComponent.tsx +++ b/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddOfflineTenantComponent.tsx @@ -148,4 +148,4 @@ export default function AddOfflineTenantComponent({ ); -} \ No newline at end of file +} diff --git a/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddPartionComponent.tsx b/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddPartionComponent.tsx index deab732021d1..51bb756eacf9 100644 --- a/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddPartionComponent.tsx +++ b/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddPartionComponent.tsx @@ -260,4 +260,4 @@ export default function AddPartionComponent({ ); -} \ No newline at end of file +} diff --git a/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddQueryComponent.tsx b/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddQueryComponent.tsx index 9270734fa7b6..8b167ebc8102 100644 --- a/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddQueryComponent.tsx +++ b/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddQueryComponent.tsx @@ -88,4 +88,4 @@ export default function AddQueryComponent({ ); -} \ No newline at end of file +} diff --git a/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddRealTimeIngestionComponent.tsx b/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddRealTimeIngestionComponent.tsx index 196ef01eee09..b4ae459d2c15 100644 --- a/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddRealTimeIngestionComponent.tsx +++ b/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddRealTimeIngestionComponent.tsx @@ -169,4 +169,4 @@ export default function AddRealTimeIngestionComponent({ ); -} \ No newline at end of file +} diff --git a/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddRealTimePartionComponent.tsx b/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddRealTimePartionComponent.tsx index a8e783012396..736b00848a90 100644 --- a/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddRealTimePartionComponent.tsx +++ b/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddRealTimePartionComponent.tsx @@ -231,4 +231,4 @@ export default function AddRealTimePartionComponent({ ); -} \ No newline at end of file +} diff --git a/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddRealtimeTableOp.tsx b/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddRealtimeTableOp.tsx index e9fe46cb72a2..f08f0441f7d0 100644 --- a/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddRealtimeTableOp.tsx +++ b/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddRealtimeTableOp.tsx @@ -462,4 +462,4 @@ const checkFields = (tableObj,fields) => { ); -} \ No newline at end of file +} diff --git a/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddSchemaOp.tsx b/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddSchemaOp.tsx index f5acf630b449..9a0057a1fec4 100644 --- a/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddSchemaOp.tsx +++ b/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddSchemaOp.tsx @@ -47,9 +47,9 @@ type Props = { }; const defaultSchemaConfig = { - schemaName:'', - dimensionFieldSpecs: [], - metricFieldSpecs: [], + schemaName:'', + dimensionFieldSpecs: [], + metricFieldSpecs: [], dateTimeFieldSpecs: [] }; @@ -258,9 +258,9 @@ export default function AddSchemaOp({ )} {editView === EditView.JSON && ( - { try{ const jsonSchema = JSON.parse(newValue); @@ -268,10 +268,10 @@ export default function AddSchemaOp({ setJsonSchema(jsonSchema); } }catch(e){} - }} + }} /> )} ); -} \ No newline at end of file +} diff --git a/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddStorageComponent.tsx b/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddStorageComponent.tsx index 3f1515ec53b3..4bb6c319ab39 100644 --- a/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddStorageComponent.tsx +++ b/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddStorageComponent.tsx @@ -115,4 +115,4 @@ export default function AddStorageComponent({ ); -} \ No newline at end of file +} diff --git a/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddTableComponent.tsx b/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddTableComponent.tsx index 6391beccab20..1e041560db2d 100644 --- a/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddTableComponent.tsx +++ b/pinot-controller/src/main/resources/app/components/Homepage/Operations/AddTableComponent.tsx @@ -124,7 +124,7 @@ export default function AddTableComponent({ onChange={(e)=> changeHandler('tableName', e.target.value)} /> - + Table Type {requiredAstrix} + {/* Search field */} + e.stopPropagation()}> + setSearchTerm(e.target.value)} + onClick={(e) => e.stopPropagation()} + onMouseDown={(e) => e.stopPropagation()} + onKeyDown={(e) => e.stopPropagation()} + onKeyUp={(e) => e.stopPropagation()} + size="small" + fullWidth + variant="outlined" + className={classes.searchField} + /> + + + + + {/* Standardized timezones section */} + + + Timezones (UTC Offsets & 3-Letter Codes) + + + {filteredTimezones.map(renderMenuItem)} + + + ); +}; + +export default TimezoneSelector; diff --git a/pinot-controller/src/main/resources/app/components/User/AddUser.tsx b/pinot-controller/src/main/resources/app/components/User/AddUser.tsx index 08002a1f58ad..7a3dced66b2f 100644 --- a/pinot-controller/src/main/resources/app/components/User/AddUser.tsx +++ b/pinot-controller/src/main/resources/app/components/User/AddUser.tsx @@ -240,4 +240,4 @@ export default function AddUser({ ); -} \ No newline at end of file +} diff --git a/pinot-controller/src/main/resources/app/components/User/UpdateUser.tsx b/pinot-controller/src/main/resources/app/components/User/UpdateUser.tsx index bdcbe29f99ec..c1a1e1c7b901 100644 --- a/pinot-controller/src/main/resources/app/components/User/UpdateUser.tsx +++ b/pinot-controller/src/main/resources/app/components/User/UpdateUser.tsx @@ -232,4 +232,4 @@ export default function UpdateUser({ ); -} \ No newline at end of file +} diff --git a/pinot-controller/src/main/resources/app/components/Zookeeper/TreeDirectory.tsx b/pinot-controller/src/main/resources/app/components/Zookeeper/TreeDirectory.tsx index e04234aab3ec..63c8ec05337f 100644 --- a/pinot-controller/src/main/resources/app/components/Zookeeper/TreeDirectory.tsx +++ b/pinot-controller/src/main/resources/app/components/Zookeeper/TreeDirectory.tsx @@ -289,4 +289,4 @@ const TreeDirectory = ({ ); }; -export default TreeDirectory; \ No newline at end of file +export default TreeDirectory; diff --git a/pinot-controller/src/main/resources/app/components/auth/AuthProvider.tsx b/pinot-controller/src/main/resources/app/components/auth/AuthProvider.tsx index c510438bc0d3..e41b79e89aab 100644 --- a/pinot-controller/src/main/resources/app/components/auth/AuthProvider.tsx +++ b/pinot-controller/src/main/resources/app/components/auth/AuthProvider.tsx @@ -92,7 +92,7 @@ export const AuthProvider = ({ children }) => { // basic auth is handled by login page } - // set OIDC auth details + // set OIDC auth details if (authWorkFlowInternal === AuthWorkflow.OIDC) { const issuer = authInfoResponse && authInfoResponse.issuer ? authInfoResponse.issuer : ''; @@ -121,7 +121,7 @@ export const AuthProvider = ({ children }) => { setAccessToken(accessToken); setAuthUserName(PinotMethodUtils.getAuthUserNameFromAccessToken(accessToken.replace("Bearer ", ""))) setAuthUserEmail(PinotMethodUtils.getAuthUserEmailFromAccessToken(accessToken.replace("Bearer ", ""))) - + initAxios(accessToken); setAuthenticated(true); diff --git a/pinot-controller/src/main/resources/app/components/useMinionMetaData.tsx b/pinot-controller/src/main/resources/app/components/useMinionMetaData.tsx index 4d5e22ffe188..5a44006c9a29 100644 --- a/pinot-controller/src/main/resources/app/components/useMinionMetaData.tsx +++ b/pinot-controller/src/main/resources/app/components/useMinionMetaData.tsx @@ -71,4 +71,4 @@ export default function useMinionMetadata(props) { ) }; -} \ No newline at end of file +} diff --git a/pinot-controller/src/main/resources/app/components/useScheduleAdhocModal.tsx b/pinot-controller/src/main/resources/app/components/useScheduleAdhocModal.tsx index 502428444708..308954528027 100644 --- a/pinot-controller/src/main/resources/app/components/useScheduleAdhocModal.tsx +++ b/pinot-controller/src/main/resources/app/components/useScheduleAdhocModal.tsx @@ -81,4 +81,4 @@ export default function useScheduleAdhocModal() { handleSheduleAdhoc, dialog }; -} \ No newline at end of file +} diff --git a/pinot-controller/src/main/resources/app/components/useTaskListing.tsx b/pinot-controller/src/main/resources/app/components/useTaskListing.tsx index 97d43f765005..5e6d6f6972d1 100644 --- a/pinot-controller/src/main/resources/app/components/useTaskListing.tsx +++ b/pinot-controller/src/main/resources/app/components/useTaskListing.tsx @@ -21,9 +21,11 @@ import React, { useEffect, useState } from 'react'; import { TableData } from 'Models'; import CustomizedTables from './Table'; import PinotMethodUtils from '../utils/PinotMethodUtils'; +import { useTimezone } from '../contexts/TimezoneContext'; export default function useTaskListing(props) { const { taskType, tableName } = props; + const { currentTimezone } = useTimezone(); const [fetching, setFetching] = useState(true); const [tasks, setTasks] = useState({ records: [], columns: [] }); @@ -36,7 +38,7 @@ export default function useTaskListing(props) { useEffect(() => { fetchData(); - }, []); + }, [currentTimezone]); return { tasks, @@ -52,4 +54,4 @@ export default function useTaskListing(props) { /> ) }; -} \ No newline at end of file +} diff --git a/pinot-controller/src/main/resources/app/contexts/TimezoneContext.tsx b/pinot-controller/src/main/resources/app/contexts/TimezoneContext.tsx new file mode 100644 index 000000000000..eeb97728f3b5 --- /dev/null +++ b/pinot-controller/src/main/resources/app/contexts/TimezoneContext.tsx @@ -0,0 +1,66 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import React, { createContext, useContext, useState, useEffect, ReactNode } from 'react'; +import { getCurrentTimezone, setCurrentTimezone, DEFAULT_TIMEZONE } from '../utils/TimezoneUtils'; + +interface TimezoneContextType { + currentTimezone: string; + setTimezone: (timezone: string) => void; +} + +const TimezoneContext = createContext(undefined); + +interface TimezoneProviderProps { + children: ReactNode; +} + +export const TimezoneProvider: React.FC = ({ children }) => { + const [currentTimezone, setCurrentTimezoneState] = useState(DEFAULT_TIMEZONE); + + useEffect(() => { + // Initialize timezone from localStorage + const savedTimezone = getCurrentTimezone(); + setCurrentTimezoneState(savedTimezone); + }, []); + + const setTimezone = (timezone: string) => { + setCurrentTimezoneState(timezone); + setCurrentTimezone(timezone); + }; + + const value: TimezoneContextType = { + currentTimezone, + setTimezone, + }; + + return ( + + {children} + + ); +}; + +export const useTimezone = (): TimezoneContextType => { + const context = useContext(TimezoneContext); + if (context === undefined) { + throw new Error('useTimezone must be used within a TimezoneProvider'); + } + return context; +}; diff --git a/pinot-controller/src/main/resources/app/index.tsx b/pinot-controller/src/main/resources/app/index.tsx index d06d57bbc491..caa5a2bec30e 100644 --- a/pinot-controller/src/main/resources/app/index.tsx +++ b/pinot-controller/src/main/resources/app/index.tsx @@ -37,6 +37,6 @@ ReactDOM.render( - , + , document.getElementById('app') ); diff --git a/pinot-controller/src/main/resources/app/pages/HomePage.tsx b/pinot-controller/src/main/resources/app/pages/HomePage.tsx index e2909b61a851..d07249acd1c5 100644 --- a/pinot-controller/src/main/resources/app/pages/HomePage.tsx +++ b/pinot-controller/src/main/resources/app/pages/HomePage.tsx @@ -208,10 +208,10 @@ const HomePage = () => { - {taskTypesTable} diff --git a/pinot-controller/src/main/resources/app/pages/InstanceListingPage.tsx b/pinot-controller/src/main/resources/app/pages/InstanceListingPage.tsx index 770142fdc104..730c73bc4dfa 100644 --- a/pinot-controller/src/main/resources/app/pages/InstanceListingPage.tsx +++ b/pinot-controller/src/main/resources/app/pages/InstanceListingPage.tsx @@ -68,14 +68,14 @@ const InstanceListingPage = () => { ) : ( - ); }; -export default InstanceListingPage; \ No newline at end of file +export default InstanceListingPage; diff --git a/pinot-controller/src/main/resources/app/pages/LoginPage.tsx b/pinot-controller/src/main/resources/app/pages/LoginPage.tsx index 37d4b3b18e5a..519b8885e482 100644 --- a/pinot-controller/src/main/resources/app/pages/LoginPage.tsx +++ b/pinot-controller/src/main/resources/app/pages/LoginPage.tsx @@ -126,4 +126,4 @@ const LoginPage = (props) => { ); }; -export default LoginPage; \ No newline at end of file +export default LoginPage; diff --git a/pinot-controller/src/main/resources/app/pages/MinionTaskManager.tsx b/pinot-controller/src/main/resources/app/pages/MinionTaskManager.tsx index bc28be624c81..7f921acbc704 100644 --- a/pinot-controller/src/main/resources/app/pages/MinionTaskManager.tsx +++ b/pinot-controller/src/main/resources/app/pages/MinionTaskManager.tsx @@ -45,4 +45,4 @@ const MinionTaskManager = () => { ); }; -export default MinionTaskManager; \ No newline at end of file +export default MinionTaskManager; diff --git a/pinot-controller/src/main/resources/app/pages/Query.tsx b/pinot-controller/src/main/resources/app/pages/Query.tsx index 868dcb7729f6..1f5caeef672d 100644 --- a/pinot-controller/src/main/resources/app/pages/Query.tsx +++ b/pinot-controller/src/main/resources/app/pages/Query.tsx @@ -298,7 +298,7 @@ const QueryPage = () => { useEffect(() => { handleQueryInterfaceKeyDownRef.current = handleQueryInterfaceKeyDown; }, [handleQueryInterfaceKeyDown]); - + const handleComment = (cm: NativeCodeMirror.Editor) => { const selections = cm.listSelections(); @@ -566,7 +566,7 @@ const QueryPage = () => { onChange={handleOutputDataChange} // Ensures the latest function is always called, preventing stale state issues due to closures. // Directly passing handleQueryInterfaceKeyDown may result in outdated state references. - onKeyDown={(editor, event) => handleQueryInterfaceKeyDownRef.current(editor, event)} + onKeyDown={(editor, event) => handleQueryInterfaceKeyDownRef.current(editor, event)} className={classes.codeMirror} autoCursor={false} /> @@ -646,13 +646,13 @@ const QueryPage = () => { ) } - + {/* Sql result errors */} {resultError && resultError.length > 0 && ( <> - { ) } - + {resultData.columns.length ? ( <> @@ -781,7 +781,7 @@ const QueryPage = () => { showSearchBox={true} inAccordionFormat={true} /> - )} + )} {resultViewType === ResultViewType.JSON && ( ) => { const classes = useStyles(); + const { currentTimezone } = useTimezone(); const history = useHistory(); const location = useLocation(); const { tableName, segmentName: encodedSegmentName } = match.params; @@ -118,7 +125,12 @@ const SegmentDetails = ({ match }: RouteComponentProps) => { const initialSummary = { segmentName, totalDocs: null, + segmentSizeInBytes: null, createTime: null, + pushTime: null, + refreshTime: null, + startTime: null, + endTime: null, }; const [segmentSummary, setSegmentSummary] = useState(initialSummary); @@ -172,12 +184,21 @@ const SegmentDetails = ({ match }: RouteComponentProps) => { const setSummary = (segmentMetadata: SegmentMetadata) => { const segmentMetaDataJson = { ...segmentMetadata }; + + // Get the time unit from segment metadata for proper time conversion + // The start/end times are stored in the segment's time unit and need to be converted to milliseconds + // Creation, push, and refresh times are always in millisecond epoch format + const timeUnit = segmentMetaDataJson['segment.time.unit'] as string; + setSegmentSummary({ segmentName, totalDocs: segmentMetaDataJson['segment.total.docs'] || 0, - createTime: moment(+segmentMetaDataJson['segment.creation.time']).format( - 'MMMM Do YYYY, h:mm:ss' - ), + segmentSizeInBytes: segmentMetaDataJson['segment.size.in.bytes'], + createTime: convertTimeToMilliseconds(segmentMetaDataJson['segment.creation.time'], 'MILLISECONDS'), + pushTime: convertTimeToMilliseconds(segmentMetaDataJson['segment.push.time'], 'MILLISECONDS'), + refreshTime: convertTimeToMilliseconds(segmentMetaDataJson['segment.refresh.time'], 'MILLISECONDS'), + startTime: convertTimeToMilliseconds(segmentMetaDataJson['segment.start.time'], timeUnit), + endTime: convertTimeToMilliseconds(segmentMetaDataJson['segment.end.time'], timeUnit), }); }; @@ -282,7 +303,7 @@ const SegmentDetails = ({ match }: RouteComponentProps) => { useEffect(() => { fetchData(); - }, []); + }, [currentTimezone]); const closeDialog = () => { setConfirmDialog(false); @@ -420,34 +441,70 @@ const SegmentDetails = ({ match }: RouteComponentProps) => {
    - + Segment Name:{' '} {unescape(segmentSummary.segmentName)} - - - Total Docs: - - - {segmentSummary.totalDocs ? ( - segmentSummary.totalDocs - ) : ( - - )} - + + Total Docs:{' '} + {segmentSummary.totalDocs ? ( + Number(segmentSummary.totalDocs).toLocaleString() + ) : ( + + )} + + + Segment Size:{' '} + {segmentSummary.segmentSizeInBytes ? ( + Utils.formatBytes(Number(segmentSummary.segmentSizeInBytes)) + ) : ( + 'N/A' + )} + + + + Create Time:{' '} + {segmentSummary.createTime && typeof segmentSummary.createTime === 'number' ? ( + formatTimeInTimezone(segmentSummary.createTime, 'MMMM Do YYYY, h:mm:ss z') + ) : segmentSummary.createTime === null ? ( + 'N/A' + ) : ( + + )} + + + Push Time:{' '} + {segmentSummary.pushTime ? ( + formatTimeInTimezone(segmentSummary.pushTime, 'MMMM Do YYYY, h:mm:ss z') + ) : ( + 'N/A' + )} + + + Refresh Time:{' '} + {segmentSummary.refreshTime ? ( + formatTimeInTimezone(segmentSummary.refreshTime, 'MMMM Do YYYY, h:mm:ss z') + ) : ( + 'N/A' + )} + + + Start Time:{' '} + {segmentSummary.startTime ? ( + formatTimeInTimezone(segmentSummary.startTime, 'MMMM Do YYYY, h:mm:ss z') + ) : ( + 'N/A' + )} - - - Create Time: - - - {segmentSummary.createTime ? ( - segmentSummary.createTime - ) : ( - - )} - + + End Time:{' '} + {segmentSummary.endTime ? ( + formatTimeInTimezone(segmentSummary.endTime, 'MMMM Do YYYY, h:mm:ss z') + ) : ( + 'N/A' + )} +
    diff --git a/pinot-controller/src/main/resources/app/pages/SubTaskDetail.tsx b/pinot-controller/src/main/resources/app/pages/SubTaskDetail.tsx index 27ab5dc27331..282e14ba1cd5 100644 --- a/pinot-controller/src/main/resources/app/pages/SubTaskDetail.tsx +++ b/pinot-controller/src/main/resources/app/pages/SubTaskDetail.tsx @@ -25,6 +25,8 @@ import SimpleAccordion from '../components/SimpleAccordion'; import PinotMethodUtils from '../utils/PinotMethodUtils'; import { TaskProgressStatus } from 'Models'; import moment from 'moment'; +import { formatTimeInTimezone } from '../utils/TimezoneUtils'; +import { useTimezone } from '../contexts/TimezoneContext'; const jsonoptions = { lineNumbers: true, @@ -69,14 +71,15 @@ const useStyles = makeStyles(() => ({ '& .CodeMirror': { maxHeight: 400 }, }, taskDetailContainer: { - overflow: "auto", - whiteSpace: "pre", + overflow: "auto", + whiteSpace: "pre", height: 600 } })); const TaskDetail = (props) => { const classes = useStyles(); + const { currentTimezone } = useTimezone(); const { subTaskID, taskID } = props.match.params; const [taskDebugData, setTaskDebugData] = useState({}); const [taskProgressData, setTaskProgressData] = useState(""); @@ -95,7 +98,7 @@ const TaskDetail = (props) => { useEffect(() => { fetchTaskDebugData(); fetchTaskProgressData(); - }, []); + }, [currentTimezone]); return ( @@ -107,12 +110,16 @@ const TaskDetail = (props) => { Status: {get(taskDebugData, 'state', '')} - - Start Time: {get(taskDebugData, 'startTime', '')} - - - Finish Time: {get(taskDebugData, 'finishTime', '')} - + {get(taskDebugData, 'startTime') && ( + + Start Time: {formatTimeInTimezone(get(taskDebugData, 'startTime'), 'MMMM Do YYYY, HH:mm:ss z')} + + )} + {get(taskDebugData, 'finishTime') && ( + + Finish Time: {formatTimeInTimezone(get(taskDebugData, 'finishTime'), 'MMMM Do YYYY, HH:mm:ss z')} + + )} Triggered By: {get(taskDebugData, 'triggeredBy', '')} @@ -165,9 +172,9 @@ const TaskDetail = (props) => { {typeof taskProgressData !== "string" && taskProgressData.map((data, index) => (
    - {data.status}} + {data.status}} /> @@ -183,4 +190,4 @@ const TaskDetail = (props) => { ); }; -export default TaskDetail; \ No newline at end of file +export default TaskDetail; diff --git a/pinot-controller/src/main/resources/app/pages/TaskDetail.tsx b/pinot-controller/src/main/resources/app/pages/TaskDetail.tsx index eff82fd3fd25..94468a3982ee 100644 --- a/pinot-controller/src/main/resources/app/pages/TaskDetail.tsx +++ b/pinot-controller/src/main/resources/app/pages/TaskDetail.tsx @@ -26,6 +26,8 @@ import { TaskRuntimeConfig } from 'Models'; import AppLoader from '../components/AppLoader'; import SimpleAccordion from '../components/SimpleAccordion'; import CustomCodemirror from '../components/CustomCodemirror'; +import { formatTimeInTimezone } from '../utils/TimezoneUtils'; +import { useTimezone } from '../contexts/TimezoneContext'; const useStyles = makeStyles(() => ({ gridContainer: { @@ -68,6 +70,7 @@ const useStyles = makeStyles(() => ({ const TaskDetail = (props) => { const classes = useStyles(); + const { currentTimezone } = useTimezone(); const { taskID, taskType, queueTableName } = props.match.params; const [fetching, setFetching] = useState(true); @@ -78,7 +81,7 @@ const TaskDetail = (props) => { const fetchData = async () => { setFetching(true); const [debugRes, runtimeConfig] = await Promise.all([ - PinotMethodUtils.getTaskDebugData(taskID), + PinotMethodUtils.getTaskDebugData(taskID), PinotMethodUtils.getTaskRuntimeConfigData(taskID) ]); const subtaskTableRecords = []; @@ -86,8 +89,8 @@ const TaskDetail = (props) => { subtaskTableRecords.push([ get(subTask, 'taskId'), get(subTask, 'state'), - get(subTask, 'startTime'), - get(subTask, 'finishTime'), + get(subTask, 'startTime') ? formatTimeInTimezone(get(subTask, 'startTime'), 'MMMM Do YYYY, HH:mm:ss z') : '-', + get(subTask, 'finishTime') ? formatTimeInTimezone(get(subTask, 'finishTime'), 'MMMM Do YYYY, HH:mm:ss z') : '-', get(subTask, 'participant'), ]) }); @@ -102,7 +105,7 @@ const TaskDetail = (props) => { useEffect(() => { fetchData(); - }, []); + }, [currentTimezone]); if(fetching) { return @@ -118,12 +121,16 @@ const TaskDetail = (props) => { Status: {get(taskDebugData, 'taskState', '')} - - Start Time: {get(taskDebugData, 'startTime', '')} - - - Finish Time: {get(taskDebugData, 'finishTime', '')} - + {get(taskDebugData, 'startTime') && ( + + Start Time: {formatTimeInTimezone(get(taskDebugData, 'startTime'), 'MMMM Do YYYY, HH:mm:ss z')} + + )} + {get(taskDebugData, 'finishTime') && ( + + Finish Time: {formatTimeInTimezone(get(taskDebugData, 'finishTime'), 'MMMM Do YYYY, HH:mm:ss z')} + + )} Triggered By: {get(taskDebugData, 'triggeredBy', '')} @@ -145,7 +152,7 @@ const TaskDetail = (props) => { /> - + {/* Sub task table */} { ); }; -export default TaskDetail; \ No newline at end of file +export default TaskDetail; diff --git a/pinot-controller/src/main/resources/app/pages/TaskQueue.tsx b/pinot-controller/src/main/resources/app/pages/TaskQueue.tsx index 3bd62ef7a235..380e9ee87b64 100644 --- a/pinot-controller/src/main/resources/app/pages/TaskQueue.tsx +++ b/pinot-controller/src/main/resources/app/pages/TaskQueue.tsx @@ -196,4 +196,4 @@ const TaskQueue = (props) => { ); }; -export default TaskQueue; \ No newline at end of file +export default TaskQueue; diff --git a/pinot-controller/src/main/resources/app/pages/TaskQueueTable.tsx b/pinot-controller/src/main/resources/app/pages/TaskQueueTable.tsx index 5d88e22140f8..6868872c268b 100644 --- a/pinot-controller/src/main/resources/app/pages/TaskQueueTable.tsx +++ b/pinot-controller/src/main/resources/app/pages/TaskQueueTable.tsx @@ -20,6 +20,8 @@ import React, { useEffect, useState, useContext } from 'react'; import { get, keys, last } from 'lodash'; import moment from 'moment'; +import { formatTimeInTimezone } from '../utils/TimezoneUtils'; +import { useTimezone } from '../contexts/TimezoneContext'; import { UnControlled as CodeMirror } from 'react-codemirror2'; import { Grid, makeStyles, Box } from '@material-ui/core'; import { NotificationContext } from '../components/Notification/NotificationContext'; @@ -78,6 +80,7 @@ const useStyles = makeStyles(() => ({ const TaskQueueTable = (props) => { const classes = useStyles(); + const { currentTimezone } = useTimezone(); const { taskType, queueTableName: tableName } = props.match.params; const {dispatch} = useContext(NotificationContext); @@ -104,7 +107,7 @@ const TaskQueueTable = (props) => { useEffect(() => { fetchData(); - }, []); + }, [currentTimezone]); const handleScheduleNow = async () => { const res = await PinotMethodUtils.scheduleTaskAction(tableName, taskType); @@ -139,7 +142,7 @@ const TaskQueueTable = (props) => { show: true }); } - + await fetchData(); scheduleNowConfirm.setConfirmDialog(false); }; @@ -189,10 +192,10 @@ const TaskQueueTable = (props) => { Cron Schedule : {cronExpression} - Previous Fire Time: {previousFireTime ? moment(previousFireTime).toString() : '-'} + Previous Fire Time: {previousFireTime ? formatTimeInTimezone(previousFireTime) : '-'} - Next Fire Time: {nextFireTime ? moment(nextFireTime).toString() : '-'} + Next Fire Time: {nextFireTime ? formatTimeInTimezone(nextFireTime) : '-'}
    @@ -237,4 +240,4 @@ const TaskQueueTable = (props) => { ); }; -export default TaskQueueTable; \ No newline at end of file +export default TaskQueueTable; diff --git a/pinot-controller/src/main/resources/app/pages/TenantDetails.tsx b/pinot-controller/src/main/resources/app/pages/TenantDetails.tsx index 37862ec5accb..8bfc649b59a0 100644 --- a/pinot-controller/src/main/resources/app/pages/TenantDetails.tsx +++ b/pinot-controller/src/main/resources/app/pages/TenantDetails.tsx @@ -557,7 +557,7 @@ const TenantPageDetails = ({ match }: RouteComponentProps) => { const result = await PinotMethodUtils.rebalanceBrokersForTableOp(tableName); syncResponse(result); }; - + const handleRepairTable = () => { setDialogDetails({ title: 'Repair Table', diff --git a/pinot-controller/src/main/resources/app/pages/ZookeeperPage.tsx b/pinot-controller/src/main/resources/app/pages/ZookeeperPage.tsx index 97b3c43c37eb..e1a62c831d00 100644 --- a/pinot-controller/src/main/resources/app/pages/ZookeeperPage.tsx +++ b/pinot-controller/src/main/resources/app/pages/ZookeeperPage.tsx @@ -95,7 +95,7 @@ const ZookeeperPage = () => { const handleChange = (event: React.ChangeEvent<{}>, newValue: number) => { setValue(newValue); }; - + // fetch and show children tree if not already fetched const showChildEvent = (pathObj) => { if(!pathObj.hasChildRendered){ diff --git a/pinot-controller/src/main/resources/app/requests/index.ts b/pinot-controller/src/main/resources/app/requests/index.ts index 4bef8fdff299..7c4d05551cdc 100644 --- a/pinot-controller/src/main/resources/app/requests/index.ts +++ b/pinot-controller/src/main/resources/app/requests/index.ts @@ -54,13 +54,13 @@ import { PauseStatusDetails } from 'Models'; +import { baseApi, baseApiWithErrors, transformApi } from '../utils/axios-config'; + const headers = { 'Content-Type': 'application/json; charset=UTF-8', 'Accept': 'text/plain, */*; q=0.01' }; -import { baseApi, baseApiWithErrors, transformApi } from '../utils/axios-config'; - export const getTenants = (): Promise> => baseApi.get('/tenants'); @@ -87,7 +87,7 @@ export const getSchema = (name: string): Promise> => { let queryParams = {}; - + if(reload) { queryParams["reload"] = reload; } @@ -198,10 +198,10 @@ export const executeTask = (data): Promise> => export const getJobDetail = (tableName: string, taskType: string): Promise> => baseApi.get(`/tasks/scheduler/jobDetails?tableName=${tableName}&taskType=${taskType}`, { headers: { ...headers, Accept: 'application/json' } }); - + export const getMinionMeta = (tableName: string, taskType: string): Promise> => baseApi.get(`/tasks/${taskType}/${tableName}/metadata`, { headers: { ...headers, Accept: 'application/json' } }); - + export const getTasks = (tableName: string, taskType: string): Promise> => baseApi.get(`/tasks/${taskType}/${tableName}/state`, { headers: { ...headers, Accept: 'application/json' } }); diff --git a/pinot-controller/src/main/resources/app/router.tsx b/pinot-controller/src/main/resources/app/router.tsx index a6c1b4483338..14175c1919a0 100644 --- a/pinot-controller/src/main/resources/app/router.tsx +++ b/pinot-controller/src/main/resources/app/router.tsx @@ -65,4 +65,4 @@ export default [ { path: '/zookeeper', Component: ZookeeperPage }, { path: '/login', Component: LoginPage }, { path: '/user', Component: UserPage} -]; \ No newline at end of file +]; diff --git a/pinot-controller/src/main/resources/app/utils/PinotMethodUtils.ts b/pinot-controller/src/main/resources/app/utils/PinotMethodUtils.ts index e29ffff3583a..e9428b3ad7d0 100644 --- a/pinot-controller/src/main/resources/app/utils/PinotMethodUtils.ts +++ b/pinot-controller/src/main/resources/app/utils/PinotMethodUtils.ts @@ -171,7 +171,7 @@ const getAllInstances = () => { [InstanceType.SERVER]: [], [InstanceType.MINION]: [] }; - + data.instances.forEach((instance) => { const instanceType = instance.split('_')[0].toUpperCase(); instanceTypeToInstancesMap[instanceType].push(instance); @@ -319,7 +319,7 @@ const getQueryResults = (params) => { } if (queryResponse && queryResponse.exceptions && queryResponse.exceptions.length) { exceptions = queryResponse.exceptions as SqlException[]; - } + } if (queryResponse.resultTable?.dataSchema?.columnNames?.length) { columnList = queryResponse.resultTable.dataSchema.columnNames; dataArray = queryResponse.resultTable.rows; @@ -571,16 +571,16 @@ const getExternalViewObj = (tableName) => { return getExternalView(tableName).then((result) => { return result.data.OFFLINE || result.data.REALTIME; }); -}; +}; const fetchServerToSegmentsCountData = (tableName, tableType) => { return getServerToSegmentsCount(tableName, tableType).then((results) => { - const segmentsArray = results.data; + const segmentsArray = results.data; return { records: segmentsArray.flatMap((server) => - Object.entries(server.serverToSegmentsCountMap).map(([serverName, segmentsCount]) => [ - serverName, - segmentsCount + Object.entries(server.serverToSegmentsCountMap).map(([serverName, segmentsCount]) => [ + serverName, + segmentsCount ]) ) }; @@ -932,6 +932,7 @@ const getElapsedTime = (startTime) => { } const getTasksList = async (tableName, taskType) => { + const { formatTimeInTimezone } = await import('./TimezoneUtils'); const finalResponse = { columns: ['Task ID', 'Status', 'Start Time', 'Finish Time', 'Num of Sub Tasks'], records: [] @@ -944,8 +945,8 @@ const getTasksList = async (tableName, taskType) => { finalResponse.records.push([ taskID, status, - get(debugData, 'data.startTime', ''), - get(debugData, 'data.finishTime', ''), + get(debugData, 'data.startTime') ? formatTimeInTimezone(get(debugData, 'data.startTime'), 'MMMM Do YYYY, HH:mm:ss z') : '', + get(debugData, 'data.finishTime') ? formatTimeInTimezone(get(debugData, 'data.finishTime'), 'MMMM Do YYYY, HH:mm:ss z') : '', get(debugData, 'data.subtaskCount.total', 0) ]); }; @@ -961,7 +962,7 @@ const getTasksList = async (tableName, taskType) => { const getTaskRuntimeConfigData = async (taskName: string) => { const response = await getTaskRuntimeConfig(taskName); - + return response.data; } @@ -1013,7 +1014,7 @@ const deleteSegmentOp = (tableName, segmentName) => { const fetchTableJobs = async (tableName: string, jobTypes?: string) => { const response = await getTableJobs(tableName, jobTypes); - + return response.data; } @@ -1031,7 +1032,7 @@ const fetchRebalanceTableJobs = async (tableName: string): Promise { const response = await getSegmentReloadStatus(jobId); - + return response.data; } @@ -1142,7 +1143,7 @@ const verifyAuth = (authToken) => { const getAccessTokenFromHashParams = () => { let accessToken = ''; const hashParam = removeAllLeadingForwardSlash(location.hash.substring(1)); - + const urlSearchParams = new URLSearchParams(hashParam); if (urlSearchParams.has('access_token')) { accessToken = urlSearchParams.get('access_token') as string; @@ -1181,7 +1182,7 @@ const validateRedirectPath = (path: string): boolean => { const knownAppRoutes = RouterData.map((data) => data.path); const routeMatches = matchPath(pathName, {path: knownAppRoutes, exact: true}); - + if(!routeMatches) { return false; } @@ -1218,7 +1219,7 @@ const getURLWithoutAccessToken = (fallbackUrl = '/'): string => { if(urlSearchParams.toString()){ urlParams.unshift(urlSearchParams.toString()); } - + url = urlParams.join('&'); if(!validateRedirectPath(url)) { diff --git a/pinot-controller/src/main/resources/app/utils/TimezoneUtils.ts b/pinot-controller/src/main/resources/app/utils/TimezoneUtils.ts new file mode 100644 index 000000000000..d1b4851366a3 --- /dev/null +++ b/pinot-controller/src/main/resources/app/utils/TimezoneUtils.ts @@ -0,0 +1,610 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import moment from 'moment'; +import 'moment-timezone'; + +// Default timezone fallback (single source of truth) +export const DEFAULT_TIMEZONE_FALLBACK = 'UTC'; +// Default timezone +export const DEFAULT_TIMEZONE = Intl.DateTimeFormat().resolvedOptions().timeZone || DEFAULT_TIMEZONE_FALLBACK; + +// Standardized timezones using 3-letter zone IDs with location information +// Each timezone has only one representation, sorted by UTC offset +export const STANDARD_TIMEZONES = [ + // UTC-12 to UTC-1 + { value: 'BIT(Pacific/Baker)', label: 'BIT (UTC-12)' }, + { value: 'SST(Pacific/Samoa)', label: 'SST (UTC-11)' }, + { value: 'HST(Pacific/Honolulu)', label: 'HST (UTC-10)' }, + { value: 'HAST(America/Adak)', label: 'HAST (UTC-10)' }, + { value: 'AKST(America/Anchorage)', label: 'AKST (UTC-9)' }, + { value: 'HADT(America/Adak)', label: 'HADT (UTC-9)' }, + { value: 'PST(America/Los_Angeles)', label: 'PST (UTC-8)' }, + { value: 'AKDT(America/Anchorage)', label: 'AKDT (UTC-8)' }, + { value: 'MST(America/Denver)', label: 'MST (UTC-7)' }, + { value: 'PDT(America/Los_Angeles)', label: 'PDT (UTC-7)' }, + { value: 'CST(America/Chicago)', label: 'CST (UTC-6)' }, + { value: 'MDT(America/Denver)', label: 'MDT (UTC-6)' }, + { value: 'EST(America/New_York)', label: 'EST (UTC-5)' }, + { value: 'CDT(America/Chicago)', label: 'CDT (UTC-5)' }, + { value: 'COT(America/Bogota)', label: 'COT (UTC-5)' }, + { value: 'PET(America/Lima)', label: 'PET (UTC-5)' }, + { value: 'AST(America/Puerto_Rico)', label: 'AST (UTC-4)' }, + { value: 'EDT(America/New_York)', label: 'EDT (UTC-4)' }, + { value: 'VET(America/Caracas)', label: 'VET (UTC-4)' }, + { value: 'BOT(America/La_Paz)', label: 'BOT (UTC-4)' }, + { value: 'ART(America/Argentina/Buenos_Aires)', label: 'ART (UTC-3)' }, + { value: 'BRT(America/Sao_Paulo)', label: 'BRT (UTC-3)' }, + { value: 'UYT(America/Montevideo)', label: 'UYT (UTC-3)' }, + { value: 'GFT(America/Cayenne)', label: 'GFT (UTC-3)' }, + { value: 'FNT(America/Noronha)', label: 'FNT (UTC-2)' }, + { value: 'AZOT(Atlantic/Azores)', label: 'AZOT (UTC-1)' }, + { value: 'CVT(Atlantic/Cape_Verde)', label: 'CVT (UTC-1)' }, + + // UTC+0 + { value: 'UTC', label: 'UTC (UTC+0)' }, + { value: 'GMT', label: 'GMT (UTC+0)' }, + { value: 'WET(Europe/London)', label: 'WET (UTC+0)' }, + + // UTC+1 to UTC+3 + { value: 'CET(Europe/Paris)', label: 'CET (UTC+1)' }, + { value: 'WEST(Europe/London)', label: 'WEST (UTC+1)' }, + { value: 'BST(Europe/London)', label: 'BST (UTC+1)' }, + { value: 'WAT(Africa/Lagos)', label: 'WAT (UTC+1)' }, + { value: 'CAT(Africa/Harare)', label: 'CAT (UTC+2)' }, + { value: 'CEST(Europe/Paris)', label: 'CEST (UTC+2)' }, + { value: 'EET(Europe/Athens)', label: 'EET (UTC+2)' }, + { value: 'SAST(Africa/Johannesburg)', label: 'SAST (UTC+2)' }, + { value: 'IST(Asia/Jerusalem)', label: 'IST (UTC+2)' }, + { value: 'EAT(Africa/Nairobi)', label: 'EAT (UTC+3)' }, + { value: 'EEST(Europe/Athens)', label: 'EEST (UTC+3)' }, + { value: 'MSK(Europe/Moscow)', label: 'MSK (UTC+3)' }, + { value: 'AST(Asia/Riyadh)', label: 'AST (UTC+3)' }, + + // UTC+4 to UTC+6 + { value: 'GST(Asia/Dubai)', label: 'GST (UTC+4)' }, + { value: 'AMT(Asia/Yerevan)', label: 'AMT (UTC+4)' }, + { value: 'AZT(Asia/Baku)', label: 'AZT (UTC+4)' }, + { value: 'PKT(Asia/Karachi)', label: 'PKT (UTC+5)' }, + { value: 'UZT(Asia/Tashkent)', label: 'UZT (UTC+5)' }, + { value: 'YEKT(Asia/Yekaterinburg)', label: 'YEKT (UTC+5)' }, + { value: 'IST(Asia/Kolkata)', label: 'IST (UTC+5:30)' }, + { value: 'NPT(Asia/Kathmandu)', label: 'NPT (UTC+5:45)' }, + { value: 'BST(Asia/Dhaka)', label: 'BST (UTC+6)' }, + { value: 'BDT(Asia/Dhaka)', label: 'BDT (UTC+6)' }, + { value: 'BTT(Asia/Thimphu)', label: 'BTT (UTC+6)' }, + + // UTC+7 to UTC+9 + { value: 'ICT(Asia/Bangkok)', label: 'ICT (UTC+7)' }, + { value: 'WIB(Asia/Jakarta)', label: 'WIB (UTC+7)' }, + { value: 'KRAT(Asia/Krasnoyarsk)', label: 'KRAT (UTC+7)' }, + { value: 'CST(Asia/Shanghai)', label: 'CST (UTC+8)' }, + { value: 'SGT(Asia/Singapore)', label: 'SGT (UTC+8)' }, + { value: 'HKT(Asia/Hong_Kong)', label: 'HKT (UTC+8)' }, + { value: 'MYT(Asia/Kuala_Lumpur)', label: 'MYT (UTC+8)' }, + { value: 'WITA(Asia/Makassar)', label: 'WITA (UTC+8)' }, + { value: 'IRKT(Asia/Irkutsk)', label: 'IRKT (UTC+8)' }, + { value: 'JST(Asia/Tokyo)', label: 'JST (UTC+9)' }, + { value: 'KST(Asia/Seoul)', label: 'KST (UTC+9)' }, + { value: 'WIT(Asia/Jayapura)', label: 'WIT (UTC+9)' }, + { value: 'YAKT(Asia/Yakutsk)', label: 'YAKT (UTC+9)' }, + + // UTC+9:30 to UTC+12 + { value: 'ACST(Australia/Adelaide)', label: 'ACST (UTC+9:30)' }, + { value: 'AEST(Australia/Sydney)', label: 'AEST (UTC+10)' }, + { value: 'AEDT(Australia/Sydney)', label: 'AEDT (UTC+11)' }, + { value: 'SBT(Pacific/Guadalcanal)', label: 'SBT (UTC+11)' }, + { value: 'NCT(Pacific/Noumea)', label: 'NCT (UTC+11)' }, + { value: 'VUT(Pacific/Efate)', label: 'VUT (UTC+11)' }, + { value: 'NZST(Pacific/Auckland)', label: 'NZST (UTC+12)' }, + { value: 'FJT(Pacific/Fiji)', label: 'FJT (UTC+12)' }, + { value: 'TVT(Pacific/Funafuti)', label: 'TVT (UTC+12)' }, + { value: 'NZDT(Pacific/Auckland)', label: 'NZDT (UTC+13)' }, + { value: 'TOT(Pacific/Tongatapu)', label: 'TOT (UTC+13)' }, + { value: 'LINT(Pacific/Kiritimati)', label: 'LINT (UTC+14)' }, +]; + +// Legacy support - map common location-based timezones to standardized format +const TIMEZONE_MAPPING: Record = { + // US Timezones + 'America/New_York': 'EST(America/New_York)', + 'America/Chicago': 'CST(America/Chicago)', + 'America/Denver': 'MST(America/Denver)', + 'America/Los_Angeles': 'PST(America/Los_Angeles)', + 'America/Anchorage': 'AKST(America/Anchorage)', + 'America/Adak': 'HAST(America/Adak)', + 'Pacific/Honolulu': 'HST(Pacific/Honolulu)', + 'America/Puerto_Rico': 'AST(America/Puerto_Rico)', + + // South American Timezones + 'America/Sao_Paulo': 'BRT(America/Sao_Paulo)', + 'America/Noronha': 'FNT(America/Noronha)', + 'America/Argentina/Buenos_Aires': 'ART(America/Argentina/Buenos_Aires)', + 'America/Montevideo': 'UYT(America/Montevideo)', + 'America/Cayenne': 'GFT(America/Cayenne)', + 'America/Bogota': 'COT(America/Bogota)', + 'America/Lima': 'PET(America/Lima)', + 'America/Caracas': 'VET(America/Caracas)', + 'America/La_Paz': 'BOT(America/La_Paz)', + + // Atlantic Timezones + 'Atlantic/Azores': 'AZOT(Atlantic/Azores)', + 'Atlantic/Cape_Verde': 'CVT(Atlantic/Cape_Verde)', + + // African Timezones + 'Africa/Lagos': 'WAT(Africa/Lagos)', + 'Africa/Harare': 'CAT(Africa/Harare)', + 'Africa/Johannesburg': 'SAST(Africa/Johannesburg)', + 'Africa/Nairobi': 'EAT(Africa/Nairobi)', + + // European Timezones + 'Europe/London': 'WET(Europe/London)', + 'Europe/Paris': 'CET(Europe/Paris)', + 'Europe/Berlin': 'CET(Europe/Paris)', + 'Europe/Athens': 'EET(Europe/Athens)', + 'Europe/Moscow': 'MSK(Europe/Moscow)', + + // Middle Eastern Timezones + 'Asia/Jerusalem': 'IST(Asia/Jerusalem)', + 'Asia/Riyadh': 'AST(Asia/Riyadh)', + 'Asia/Dubai': 'GST(Asia/Dubai)', + + // Central Asian Timezones + 'Asia/Yerevan': 'AMT(Asia/Yerevan)', + 'Asia/Baku': 'AZT(Asia/Baku)', + 'Asia/Karachi': 'PKT(Asia/Karachi)', + 'Asia/Tashkent': 'UZT(Asia/Tashkent)', + 'Asia/Yekaterinburg': 'YEKT(Asia/Yekaterinburg)', + 'Asia/Kolkata': 'IST(Asia/Kolkata)', + 'Asia/Kathmandu': 'NPT(Asia/Kathmandu)', + 'Asia/Dhaka': 'BDT(Asia/Dhaka)', + 'Asia/Thimphu': 'BTT(Asia/Thimphu)', + + // Southeast Asian Timezones + 'Asia/Bangkok': 'ICT(Asia/Bangkok)', + 'Asia/Jakarta': 'WIB(Asia/Jakarta)', + 'Asia/Makassar': 'WITA(Asia/Makassar)', + 'Asia/Jayapura': 'WIT(Asia/Jayapura)', + 'Asia/Krasnoyarsk': 'KRAT(Asia/Krasnoyarsk)', + 'Asia/Irkutsk': 'IRKT(Asia/Irkutsk)', + 'Asia/Yakutsk': 'YAKT(Asia/Yakutsk)', + + // East Asian Timezones + 'Asia/Shanghai': 'CST(Asia/Shanghai)', + 'Asia/Singapore': 'SGT(Asia/Singapore)', + 'Asia/Hong_Kong': 'HKT(Asia/Hong_Kong)', + 'Asia/Kuala_Lumpur': 'MYT(Asia/Kuala_Lumpur)', + 'Asia/Tokyo': 'JST(Asia/Tokyo)', + 'Asia/Seoul': 'KST(Asia/Seoul)', + + // Pacific Timezones + 'Pacific/Baker': 'BIT(Pacific/Baker)', + 'Pacific/Samoa': 'SST(Pacific/Samoa)', + 'Pacific/Guadalcanal': 'SBT(Pacific/Guadalcanal)', + 'Pacific/Noumea': 'NCT(Pacific/Noumea)', + 'Pacific/Efate': 'VUT(Pacific/Efate)', + 'Pacific/Auckland': 'NZST(Pacific/Auckland)', + 'Pacific/Fiji': 'FJT(Pacific/Fiji)', + 'Pacific/Funafuti': 'TVT(Pacific/Funafuti)', + 'Pacific/Tongatapu': 'TOT(Pacific/Tongatapu)', + 'Pacific/Kiritimati': 'LINT(Pacific/Kiritimati)', + + // Australian Timezones + 'Australia/Sydney': 'AEST(Australia/Sydney)', + 'Australia/Adelaide': 'ACST(Australia/Adelaide)', + 'Australia/Perth': 'CST(Asia/Shanghai)', + + // Common Timezone Abbreviations + 'EST': 'EST(America/New_York)', + 'CST': 'CST(America/Chicago)', + 'MST': 'MST(America/Denver)', + 'PST': 'PST(America/Los_Angeles)', + 'EDT': 'EDT(America/New_York)', + 'CDT': 'CDT(America/Chicago)', + 'MDT': 'MDT(America/Denver)', + 'PDT': 'PDT(America/Los_Angeles)', + 'AKST': 'AKST(America/Anchorage)', + 'AKDT': 'AKDT(America/Anchorage)', + 'HAST': 'HAST(America/Adak)', + 'HADT': 'HADT(America/Adak)', + 'HST': 'HST(Pacific/Honolulu)', + 'SST': 'SST(Pacific/Samoa)', + 'AST': 'AST(America/Puerto_Rico)', + 'WET': 'WET(Europe/London)', + 'WEST': 'WEST(Europe/London)', + 'BST': 'BST(Europe/London)', + 'CET': 'CET(Europe/Paris)', + 'CEST': 'CEST(Europe/Paris)', + 'EET': 'EET(Europe/Athens)', + 'EEST': 'EEST(Europe/Athens)', + 'IST': 'IST(Asia/Kolkata)', + 'NPT': 'NPT(Asia/Kathmandu)', + 'JST': 'JST(Asia/Tokyo)', + 'KST': 'KST(Asia/Seoul)', + 'AEST': 'AEST(Australia/Sydney)', + 'AEDT': 'AEDT(Australia/Sydney)', + 'ACST': 'ACST(Australia/Adelaide)', + 'BRT': 'BRT(America/Sao_Paulo)', + 'ART': 'ART(America/Argentina/Buenos_Aires)', + 'UYT': 'UYT(America/Montevideo)', + 'GFT': 'GFT(America/Cayenne)', + 'FNT': 'FNT(America/Noronha)', + 'AZOT': 'AZOT(Atlantic/Azores)', + 'CVT': 'CVT(Atlantic/Cape_Verde)', + 'WAT': 'WAT(Africa/Lagos)', + 'CAT': 'CAT(Africa/Harare)', + 'SAST': 'SAST(Africa/Johannesburg)', + 'EAT': 'EAT(Africa/Nairobi)', + 'MSK': 'MSK(Europe/Moscow)', + 'AMT': 'AMT(Asia/Yerevan)', + 'AZT': 'AZT(Asia/Baku)', + 'GST': 'GST(Asia/Dubai)', + 'PKT': 'PKT(Asia/Karachi)', + 'UZT': 'UZT(Asia/Tashkent)', + 'YEKT': 'YEKT(Asia/Yekaterinburg)', + 'BDT': 'BDT(Asia/Dhaka)', + 'BTT': 'BTT(Asia/Thimphu)', + 'ICT': 'ICT(Asia/Bangkok)', + 'WIB': 'WIB(Asia/Jakarta)', + 'WITA': 'WITA(Asia/Makassar)', + 'WIT': 'WIT(Asia/Jayapura)', + 'KRAT': 'KRAT(Asia/Krasnoyarsk)', + 'IRKT': 'IRKT(Asia/Irkutsk)', + 'YAKT': 'YAKT(Asia/Yakutsk)', + 'SGT': 'SGT(Asia/Singapore)', + 'HKT': 'HKT(Asia/Hong_Kong)', + 'MYT': 'MYT(Asia/Kuala_Lumpur)', + 'SBT': 'SBT(Pacific/Guadalcanal)', + 'NCT': 'NCT(Pacific/Noumea)', + 'VUT': 'VUT(Pacific/Efate)', + 'BIT': 'BIT(Pacific/Baker)', + 'NZST': 'NZST(Pacific/Auckland)', + 'NZDT': 'NZDT(Pacific/Auckland)', + 'FJT': 'FJT(Pacific/Fiji)', + 'TVT': 'TVT(Pacific/Funafuti)', + 'TOT': 'TOT(Pacific/Tongatapu)', + 'LINT': 'LINT(Pacific/Kiritimati)', + 'COT': 'COT(America/Bogota)', + 'PET': 'PET(America/Lima)', + 'VET': 'VET(America/Caracas)', + 'BOT': 'BOT(America/La_Paz)', +}; + +// Time unit conversion constants (matching Java TimeUnit enum) +export const TIME_UNITS = { + NANOSECONDS: 'NANOSECONDS', + MICROSECONDS: 'MICROSECONDS', + MILLISECONDS: 'MILLISECONDS', + SECONDS: 'SECONDS', + MINUTES: 'MINUTES', + HOURS: 'HOURS', + DAYS: 'DAYS' +} as const; + +export type TimeUnit = typeof TIME_UNITS[keyof typeof TIME_UNITS]; + +// Conversion factors to milliseconds +const TIME_UNIT_TO_MS: Record = { + [TIME_UNITS.NANOSECONDS]: 1 / 1000000, + [TIME_UNITS.MICROSECONDS]: 1 / 1000, + [TIME_UNITS.MILLISECONDS]: 1, + [TIME_UNITS.SECONDS]: 1000, + [TIME_UNITS.MINUTES]: 60 * 1000, + [TIME_UNITS.HOURS]: 60 * 60 * 1000, + [TIME_UNITS.DAYS]: 24 * 60 * 60 * 1000 +}; + +/** + * Converts a time value from its original time unit to milliseconds + * + * @param timeValue - The time value to convert + * @param timeUnit - The time unit of the input value (e.g., 'SECONDS', 'DAYS') + * @returns The time value converted to milliseconds, or undefined if conversion fails + */ +export const convertTimeToMilliseconds = (timeValue: unknown, timeUnit?: string): number | undefined => { + if (timeValue === null || timeValue === undefined || timeValue === '') { + return undefined; + } + + const num = Number(timeValue); + if (isNaN(num)) { + return undefined; + } + + // If no time unit specified or unknown time unit, assume milliseconds + if (!timeUnit || !TIME_UNIT_TO_MS[timeUnit as TimeUnit]) { + return num; + } + + return num * TIME_UNIT_TO_MS[timeUnit as TimeUnit]; +}; + +/** + * Converts a timezone to its standardized format (UTC offset or 3-letter shortcut with location) + * @param timezone - The timezone to standardize + * @returns The standardized timezone string + */ +export const standardizeTimezone = (timezone: string): string => { + // If it's already a standard format, return as is + if (timezone.startsWith('UTC') || isStandardTimezone(timezone)) { + return timezone; + } + + // Map from location-based timezone to standardized format + const mapped = TIMEZONE_MAPPING[timezone]; + if (mapped) { + return mapped; + } + + // For unknown timezones, try to get the offset from moment + try { + const offset = moment.tz(timezone).format('Z'); + if (offset === 'Z') { + return 'UTC'; + } + return `UTC${offset}`; + } catch { + // Fallback to UTC if conversion fails + return 'UTC'; + } +}; + +/** + * Checks if a timezone is already in standard format + * @param timezone - The timezone to check + * @returns True if it's a standard timezone + */ +const isStandardTimezone = (timezone: string): boolean => { + return STANDARD_TIMEZONES.some(tz => tz.value === timezone); +}; + +/** + * Get standardized timezone options (UTC offsets and 3-letter shortcuts) + * @returns Array of standardized timezone options + */ +export const getStandardTimezones = () => { + return STANDARD_TIMEZONES; +}; + +// Get current timezone from localStorage or default +export const getCurrentTimezone = (): string => { + const saved = localStorage.getItem('pinot_ui:timezone'); + if (saved) { + return standardizeTimezone(saved); + } + return standardizeTimezone(DEFAULT_TIMEZONE); +}; + +// Set timezone in localStorage +export const setCurrentTimezone = (timezone: string): void => { + const standardized = standardizeTimezone(timezone); + localStorage.setItem('pinot_ui:timezone', standardized); +}; + +// Extract timezone abbreviation from standardized timezone format +const extractTimezoneAbbreviation = (standardizedTimezone: string): string => { + // For UTC, return as is + if (standardizedTimezone === 'UTC') { + return 'UTC'; + } + + // For format like 'PST(America/Los_Angeles)', extract 'PST' + const match = standardizedTimezone.match(/^([A-Z]{3,4})\(/); + if (match) { + return match[1]; + } + + // For UTC offsets like 'UTC-8', return as is + if (standardizedTimezone.startsWith('UTC')) { + return standardizedTimezone; + } + + // Fallback + return standardizedTimezone; +}; + +// Format time in the current timezone +export const formatTimeInTimezone = (time: number | string | Date, format?: string, timezone?: string): string => { + const currentTimezone = timezone || getCurrentTimezone(); + const standardizedTimezone = standardizeTimezone(currentTimezone); + + // Convert standardized timezone to moment-compatible format + const momentTimezone = convertToMomentTimezone(standardizedTimezone); + const momentTime = moment(time); + + if (momentTime.isValid()) { + const defaultFormat = 'MMMM Do YYYY, HH:mm:ss'; + const requestedFormat = format || defaultFormat; + + // Check if format includes timezone (z token) + if (requestedFormat.includes(' z')) { + // Extract the selected timezone abbreviation for consistency + const selectedAbbreviation = extractTimezoneAbbreviation(standardizedTimezone); + // Format without timezone and manually append the selected abbreviation + const formatWithoutTz = requestedFormat.replace(' z', ''); + const formattedTime = momentTime.tz(momentTimezone).format(formatWithoutTz); + return `${formattedTime} ${selectedAbbreviation}`; + } else { + // No timezone requested, format normally + return momentTime.tz(momentTimezone).format(requestedFormat); + } + } + + return 'Invalid time'; +}; + +/** + * Converts standardized timezone format to moment-compatible timezone + * @param standardizedTimezone - Timezone in standardized format (e.g., 'EST(America/New_York)', 'UTC-5') + * @returns Moment-compatible timezone string + */ +const convertToMomentTimezone = (standardizedTimezone: string): string => { + if (standardizedTimezone === 'UTC') { + return 'UTC'; + } + + // Handle 3-letter shortcuts with location information + const shortcutWithLocationMatch = standardizedTimezone.match(/^([A-Z]{3,4})\(([^)]+)\)$/); + if (shortcutWithLocationMatch) { + const [, shortcut, location] = shortcutWithLocationMatch; + return location; // Use the location part for moment + } + + // Handle 3-letter shortcuts without location (legacy support) + const shortcutMapping: Record = { + 'GMT': 'UTC', + 'BIT': 'Pacific/Baker', + 'SST': 'Pacific/Samoa', + 'HST': 'Pacific/Honolulu', + 'HAST': 'America/Adak', + 'HADT': 'America/Adak', + 'AKDT': 'America/Anchorage', + 'AKST': 'America/Anchorage', + 'PST': 'America/Los_Angeles', + 'PDT': 'America/Los_Angeles', + 'MST': 'America/Denver', + 'MDT': 'America/Denver', + 'CST': 'America/Chicago', + 'CDT': 'America/Chicago', + 'EST': 'America/New_York', + 'EDT': 'America/New_York', + 'AST': 'America/Puerto_Rico', + 'COT': 'America/Bogota', + 'PET': 'America/Lima', + 'VET': 'America/Caracas', + 'BOT': 'America/La_Paz', + 'ART': 'America/Argentina/Buenos_Aires', + 'BRT': 'America/Sao_Paulo', + 'UYT': 'America/Montevideo', + 'GFT': 'America/Cayenne', + 'FNT': 'America/Noronha', + 'AZOT': 'Atlantic/Azores', + 'CVT': 'Atlantic/Cape_Verde', + 'WET': 'Europe/London', + 'WEST': 'Europe/London', + 'BST': 'Europe/London', + 'WAT': 'Africa/Lagos', + 'CET': 'Europe/Paris', + 'CEST': 'Europe/Paris', + 'CAT': 'Africa/Harare', + 'EET': 'Europe/Athens', + 'EEST': 'Europe/Athens', + 'SAST': 'Africa/Johannesburg', + 'EAT': 'Africa/Nairobi', + 'MSK': 'Europe/Moscow', + 'IST': 'Asia/Kolkata', + 'NPT': 'Asia/Kathmandu', + 'BDT': 'Asia/Dhaka', + 'BTT': 'Asia/Thimphu', + 'AMT': 'Asia/Yerevan', + 'AZT': 'Asia/Baku', + 'GST': 'Asia/Dubai', + 'PKT': 'Asia/Karachi', + 'UZT': 'Asia/Tashkent', + 'YEKT': 'Asia/Yekaterinburg', + 'ICT': 'Asia/Bangkok', + 'WIB': 'Asia/Jakarta', + 'WITA': 'Asia/Makassar', + 'WIT': 'Asia/Jayapura', + 'KRAT': 'Asia/Krasnoyarsk', + 'IRKT': 'Asia/Irkutsk', + 'YAKT': 'Asia/Yakutsk', + 'SGT': 'Asia/Singapore', + 'HKT': 'Asia/Hong_Kong', + 'MYT': 'Asia/Kuala_Lumpur', + 'JST': 'Asia/Tokyo', + 'KST': 'Asia/Seoul', + 'ACST': 'Australia/Adelaide', + 'AEST': 'Australia/Sydney', + 'AEDT': 'Australia/Sydney', + 'SBT': 'Pacific/Guadalcanal', + 'NCT': 'Pacific/Noumea', + 'VUT': 'Pacific/Efate', + 'NZST': 'Pacific/Auckland', + 'NZDT': 'Pacific/Auckland', + 'FJT': 'Pacific/Fiji', + 'TVT': 'Pacific/Funafuti', + 'TOT': 'Pacific/Tongatapu', + 'LINT': 'Pacific/Kiritimati', + }; + + const mapped = shortcutMapping[standardizedTimezone]; + if (mapped) { + return mapped; + } + + // Handle UTC offset format + if (standardizedTimezone.startsWith('UTC')) { + const offsetMatch = standardizedTimezone.match(/^UTC([+-]\d+(?::\d+)?)$/); + if (offsetMatch) { + const offset = offsetMatch[1]; + // Convert to moment format + if (offset.includes(':')) { + return `Etc/GMT${offset.startsWith('+') ? '-' : '+'}${offset.substring(1)}`; + } else { + const hours = parseInt(offset); + const sign = hours >= 0 ? '-' : '+'; + const absHours = Math.abs(hours); + return `Etc/GMT${sign}${absHours}`; + } + } + } + + return 'UTC'; +}; + +// Convert time to a specific timezone +export const convertTimeToTimezone = (time: number | string | Date, targetTimezone: string): moment.Moment => { + const standardizedTimezone = standardizeTimezone(targetTimezone); + const momentTimezone = convertToMomentTimezone(standardizedTimezone); + return moment(time).tz(momentTimezone); +}; + +// Get timezone offset string +export const getTimezoneOffset = (timezone: string): string => { + const standardizedTimezone = standardizeTimezone(timezone); + const momentTimezone = convertToMomentTimezone(standardizedTimezone); + return moment.tz(momentTimezone).format('Z'); +}; + +// Get timezone display name +export const getTimezoneDisplayName = (timezone: string): string => { + const standardizedTimezone = standardizeTimezone(timezone); + const offset = getTimezoneOffset(standardizedTimezone); + return `${standardizedTimezone} (${offset})`; +}; + +/** + * Checks if the provided timezone string is a valid standardized timezone. + * + * @param timezone - The timezone string to validate + * @returns {boolean} True if the timezone is valid, false otherwise. + */ +export const isValidTimezone = (timezone: string): boolean => { + const standardized = standardizeTimezone(timezone); + return STANDARD_TIMEZONES.some(tz => tz.value === standardized); +}; + +// Legacy function for backward compatibility - now returns standardized timezones +export const getAllTimezones = () => { + return getStandardTimezones(); +}; + +// Legacy function for backward compatibility +export const COMMON_TIMEZONES = STANDARD_TIMEZONES; diff --git a/pinot-controller/src/main/resources/app/utils/Utils.tsx b/pinot-controller/src/main/resources/app/utils/Utils.tsx index 50ef179b8580..d115c96dfe81 100644 --- a/pinot-controller/src/main/resources/app/utils/Utils.tsx +++ b/pinot-controller/src/main/resources/app/utils/Utils.tsx @@ -33,6 +33,7 @@ import { import Loading from '../components/Loading'; import moment from "moment"; import {RebalanceServerOption} from "../components/Homepage/Operations/RebalanceServer/RebalanceServerOptions"; +import { formatTimeInTimezone } from './TimezoneUtils'; const getRebalanceConfigValue = ( rebalanceConfig: { [optionName: string]: string | boolean | number }, @@ -434,7 +435,7 @@ export const getDisplaySegmentStatus = (idealState, externalView): DISPLAY_SEGME return DISPLAY_SEGMENT_STATUS.UPDATING; } - // does not match any condition -> assume PARTIAL state as we are waiting for segments to converge + // does not match any condition -> assume PARTIAL state as we are waiting for segments to converge return DISPLAY_SEGMENT_STATUS.UPDATING; } @@ -476,7 +477,7 @@ const getLoadingTableData = (columns: string[]): TableData => { } const formatTime = (time: number, format?: string): string => { - return moment(time).format(format ?? "MMMM Do YYYY, HH:mm:ss") + return formatTimeInTimezone(time, format); } export default { diff --git a/pinot-controller/src/main/resources/package-lock.json b/pinot-controller/src/main/resources/package-lock.json index b38702a64c72..df04fd4ecd8d 100644 --- a/pinot-controller/src/main/resources/package-lock.json +++ b/pinot-controller/src/main/resources/package-lock.json @@ -34,6 +34,7 @@ "jwt-decode": "^3.1.2", "lodash": "4.17.21", "moment": "2.29.4", + "moment-timezone": "^0.5.42", "prop-types": "15.8.1", "re-resizable": "6.9.9", "react": "16.13.1", @@ -7679,6 +7680,18 @@ "node": "*" } }, + "node_modules/moment-timezone": { + "version": "0.5.42", + "resolved": "https://registry.npmjs.org/moment-timezone/-/moment-timezone-0.5.42.tgz", + "integrity": "sha512-tjI9goqwzkflKSTxJo+jC/W8riTFwEjjunssmFvAWlvNVApjbkJM7UHggyKO0q1Fd/kZVKY77H7C9A0XKhhAFw==", + "license": "MIT", + "dependencies": { + "moment": "^2.29.4" + }, + "engines": { + "node": "*" + } + }, "node_modules/move-concurrently": { "version": "1.0.1", "resolved": "https://registry.npmjs.org/move-concurrently/-/move-concurrently-1.0.1.tgz", @@ -18827,6 +18840,14 @@ "resolved": "https://registry.npmjs.org/moment/-/moment-2.29.4.tgz", "integrity": "sha512-5LC9SOxjSc2HF6vO2CyuTDNivEdoz2IvyJJGj6X8DJ0eFyfszE0QiEd+iXmBvUP3WHxSjFH/vIsA0EN00cgr8w==" }, + "moment-timezone": { + "version": "0.5.42", + "resolved": "https://registry.npmjs.org/moment-timezone/-/moment-timezone-0.5.42.tgz", + "integrity": "sha512-tjI9goqwzkflKSTxJo+jC/W8riTFwEjjunssmFvAWlvNVApjbkJM7UHggyKO0q1Fd/kZVKY77H7C9A0XKhhAFw==", + "requires": { + "moment": "^2.29.4" + } + }, "move-concurrently": { "version": "1.0.1", "resolved": "https://registry.npmjs.org/move-concurrently/-/move-concurrently-1.0.1.tgz", diff --git a/pinot-controller/src/main/resources/package.json b/pinot-controller/src/main/resources/package.json index f1189f9a7445..c4643cc261bb 100644 --- a/pinot-controller/src/main/resources/package.json +++ b/pinot-controller/src/main/resources/package.json @@ -23,9 +23,9 @@ ] }, "devDependencies": { + "@types/minimatch": "^6.0.0", "@typescript-eslint/eslint-plugin": "2.30.0", "@typescript-eslint/parser": "2.30.0", - "@types/minimatch": "^6.0.0", "assert": "^2.1.0", "clean-webpack-plugin": "3.0.0", "copy-webpack-plugin": "5.1.1", @@ -89,6 +89,7 @@ "jwt-decode": "^3.1.2", "lodash": "4.17.21", "moment": "2.29.4", + "moment-timezone": "^0.5.42", "prop-types": "15.8.1", "re-resizable": "6.9.9", "react": "16.13.1", From 1a5e0e5d3ab8639544d5bcd40e2b324d739df3bd Mon Sep 17 00:00:00 2001 From: Shounak kulkarni Date: Tue, 5 Aug 2025 20:04:12 +0530 Subject: [PATCH 076/167] Fail task cleanup validation only upon active subtasks (#16493) --- .../controller/api/resources/PinotTableRestletResource.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTableRestletResource.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTableRestletResource.java index 7f1dae48634a..5a3516595957 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTableRestletResource.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTableRestletResource.java @@ -524,7 +524,8 @@ public static void tableTasksCleanup(String tableWithType, boolean ignoreActiveT continue; } for (String taskName : taskStates.keySet()) { - if (TaskState.IN_PROGRESS.equals(taskStates.get(taskName))) { + if (TaskState.IN_PROGRESS.equals(taskStates.get(taskName)) + && pinotHelixTaskResourceManager.getTaskCount(taskName).getRunning() > 0) { pendingTasks.add(taskName); } else { pinotHelixTaskResourceManager.deleteTask(taskName, true); From 123dbfe2b02d8f098268946119cb89b026113f3e Mon Sep 17 00:00:00 2001 From: Shaurya Chaturvedi Date: Tue, 5 Aug 2025 09:47:46 -0700 Subject: [PATCH 077/167] [refactor] Switching to RoutingManager for broker request handlers (#16442) Co-authored-by: Shaurya Chaturvedi --- .../requesthandler/BaseBrokerRequestHandler.java | 6 +++--- .../BaseSingleStageBrokerRequestHandler.java | 12 +++++++----- .../requesthandler/GrpcBrokerRequestHandler.java | 4 ++-- .../MultiStageBrokerRequestHandler.java | 4 ++-- .../SingleConnectionBrokerRequestHandler.java | 4 ++-- 5 files changed, 16 insertions(+), 14 deletions(-) diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java index b4aafbea3f23..97b87ce26f7c 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java @@ -42,7 +42,6 @@ import org.apache.pinot.broker.broker.AccessControlFactory; import org.apache.pinot.broker.querylog.QueryLogger; import org.apache.pinot.broker.queryquota.QueryQuotaManager; -import org.apache.pinot.broker.routing.BrokerRoutingManager; import org.apache.pinot.common.config.provider.TableCache; import org.apache.pinot.common.metrics.BrokerMeter; import org.apache.pinot.common.metrics.BrokerMetrics; @@ -53,6 +52,7 @@ import org.apache.pinot.common.utils.request.RequestUtils; import org.apache.pinot.core.auth.Actions; import org.apache.pinot.core.auth.TargetType; +import org.apache.pinot.core.routing.RoutingManager; import org.apache.pinot.spi.accounting.ThreadResourceUsageAccountant; import org.apache.pinot.spi.auth.AuthorizationResult; import org.apache.pinot.spi.auth.TableAuthorizationResult; @@ -77,7 +77,7 @@ public abstract class BaseBrokerRequestHandler implements BrokerRequestHandler { private static final Logger LOGGER = LoggerFactory.getLogger(BaseBrokerRequestHandler.class); protected final PinotConfiguration _config; protected final String _brokerId; - protected final BrokerRoutingManager _routingManager; + protected final RoutingManager _routingManager; protected final AccessControlFactory _accessControlFactory; protected final QueryQuotaManager _queryQuotaManager; protected final TableCache _tableCache; @@ -101,7 +101,7 @@ public abstract class BaseBrokerRequestHandler implements BrokerRequestHandler { */ protected final Map _clientQueryIds; - public BaseBrokerRequestHandler(PinotConfiguration config, String brokerId, BrokerRoutingManager routingManager, + public BaseBrokerRequestHandler(PinotConfiguration config, String brokerId, RoutingManager routingManager, AccessControlFactory accessControlFactory, QueryQuotaManager queryQuotaManager, TableCache tableCache, ThreadResourceUsageAccountant resourceUsageAccountant) { _config = config; diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java index 4f409e8ad302..6aac44b91cb9 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java @@ -57,7 +57,6 @@ import org.apache.pinot.broker.broker.AccessControlFactory; import org.apache.pinot.broker.querylog.QueryLogger; import org.apache.pinot.broker.queryquota.QueryQuotaManager; -import org.apache.pinot.broker.routing.BrokerRoutingManager; import org.apache.pinot.common.config.provider.TableCache; import org.apache.pinot.common.http.MultiHttpRequest; import org.apache.pinot.common.http.MultiHttpRequestResponse; @@ -92,6 +91,7 @@ import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils; import org.apache.pinot.core.routing.ImplicitHybridTableRouteProvider; import org.apache.pinot.core.routing.LogicalTableRouteProvider; +import org.apache.pinot.core.routing.RoutingManager; import org.apache.pinot.core.routing.TableRouteInfo; import org.apache.pinot.core.routing.TableRouteProvider; import org.apache.pinot.core.routing.timeboundary.TimeBoundaryInfo; @@ -170,7 +170,7 @@ public abstract class BaseSingleStageBrokerRequestHandler extends BaseBrokerRequ protected LogicalTableRouteProvider _logicalTableRouteProvider; public BaseSingleStageBrokerRequestHandler(PinotConfiguration config, String brokerId, - BrokerRoutingManager routingManager, AccessControlFactory accessControlFactory, + RoutingManager routingManager, AccessControlFactory accessControlFactory, QueryQuotaManager queryQuotaManager, TableCache tableCache, ThreadResourceUsageAccountant accountant) { super(config, brokerId, routingManager, accessControlFactory, queryQuotaManager, tableCache, accountant); _disableGroovy = _config.getProperty(Broker.DISABLE_GROOVY, Broker.DEFAULT_DISABLE_GROOVY); @@ -1945,10 +1945,12 @@ private long setQueryTimeout(String tableNameWithType, Long logicalTableQueryTim // Use query-level timeout if exists queryTimeoutMs = queryLevelTimeoutMs; } else { - Long tableLevelTimeoutMs = _routingManager.getQueryTimeoutMs(tableNameWithType); - if (tableLevelTimeoutMs != null) { + TableConfig tableConfig = _tableCache.getTableConfig(tableNameWithType); + if (tableConfig != null + && tableConfig.getQueryConfig() != null + && tableConfig.getQueryConfig().getTimeoutMs() != null) { // Use table-level timeout if exists - queryTimeoutMs = tableLevelTimeoutMs; + queryTimeoutMs = tableConfig.getQueryConfig().getTimeoutMs(); } else if (logicalTableQueryTimeout != null) { queryTimeoutMs = logicalTableQueryTimeout; } else { diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java index 995b92d0ec73..402dbeade84c 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/GrpcBrokerRequestHandler.java @@ -28,7 +28,6 @@ import javax.annotation.concurrent.ThreadSafe; import org.apache.pinot.broker.broker.AccessControlFactory; import org.apache.pinot.broker.queryquota.QueryQuotaManager; -import org.apache.pinot.broker.routing.BrokerRoutingManager; import org.apache.pinot.common.config.GrpcConfig; import org.apache.pinot.common.config.provider.TableCache; import org.apache.pinot.common.failuredetector.FailureDetector; @@ -38,6 +37,7 @@ import org.apache.pinot.common.utils.grpc.ServerGrpcQueryClient; import org.apache.pinot.common.utils.grpc.ServerGrpcRequestBuilder; import org.apache.pinot.core.query.reduce.StreamingReduceService; +import org.apache.pinot.core.routing.RoutingManager; import org.apache.pinot.core.routing.SegmentsToQuery; import org.apache.pinot.core.routing.TableRouteInfo; import org.apache.pinot.core.transport.ServerInstance; @@ -63,7 +63,7 @@ public class GrpcBrokerRequestHandler extends BaseSingleStageBrokerRequestHandle private final FailureDetector _failureDetector; // TODO: Support TLS - public GrpcBrokerRequestHandler(PinotConfiguration config, String brokerId, BrokerRoutingManager routingManager, + public GrpcBrokerRequestHandler(PinotConfiguration config, String brokerId, RoutingManager routingManager, AccessControlFactory accessControlFactory, QueryQuotaManager queryQuotaManager, TableCache tableCache, FailureDetector failureDetector, ThreadResourceUsageAccountant accountant) { super(config, brokerId, routingManager, accessControlFactory, queryQuotaManager, tableCache, accountant); diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java index 1fc60da9d286..f62d48b9554a 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/MultiStageBrokerRequestHandler.java @@ -51,7 +51,6 @@ import org.apache.pinot.broker.broker.AccessControlFactory; import org.apache.pinot.broker.querylog.QueryLogger; import org.apache.pinot.broker.queryquota.QueryQuotaManager; -import org.apache.pinot.broker.routing.BrokerRoutingManager; import org.apache.pinot.common.config.TlsConfig; import org.apache.pinot.common.config.provider.TableCache; import org.apache.pinot.common.failuredetector.FailureDetector; @@ -71,6 +70,7 @@ import org.apache.pinot.common.utils.Timer; import org.apache.pinot.common.utils.config.QueryOptionsUtils; import org.apache.pinot.common.utils.tls.TlsUtils; +import org.apache.pinot.core.routing.RoutingManager; import org.apache.pinot.core.transport.ServerInstance; import org.apache.pinot.query.ImmutableQueryEnvironment; import org.apache.pinot.query.QueryEnvironment; @@ -132,7 +132,7 @@ public class MultiStageBrokerRequestHandler extends BaseBrokerRequestHandler { private final ExecutorService _queryCompileExecutor; protected final long _extraPassiveTimeoutMs; - public MultiStageBrokerRequestHandler(PinotConfiguration config, String brokerId, BrokerRoutingManager routingManager, + public MultiStageBrokerRequestHandler(PinotConfiguration config, String brokerId, RoutingManager routingManager, AccessControlFactory accessControlFactory, QueryQuotaManager queryQuotaManager, TableCache tableCache, MultiStageQueryThrottler queryThrottler, FailureDetector failureDetector, ThreadResourceUsageAccountant accountant) { diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java index ebbab31b3bee..e8eaf3894456 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/SingleConnectionBrokerRequestHandler.java @@ -26,7 +26,6 @@ import javax.annotation.concurrent.ThreadSafe; import org.apache.pinot.broker.broker.AccessControlFactory; import org.apache.pinot.broker.queryquota.QueryQuotaManager; -import org.apache.pinot.broker.routing.BrokerRoutingManager; import org.apache.pinot.common.config.NettyConfig; import org.apache.pinot.common.config.TlsConfig; import org.apache.pinot.common.config.provider.TableCache; @@ -39,6 +38,7 @@ import org.apache.pinot.common.response.broker.QueryProcessingException; import org.apache.pinot.common.utils.config.QueryOptionsUtils; import org.apache.pinot.core.query.reduce.BrokerReduceService; +import org.apache.pinot.core.routing.RoutingManager; import org.apache.pinot.core.routing.TableRouteInfo; import org.apache.pinot.core.transport.AsyncQueryResponse; import org.apache.pinot.core.transport.QueryResponse; @@ -70,7 +70,7 @@ public class SingleConnectionBrokerRequestHandler extends BaseSingleStageBrokerR private final FailureDetector _failureDetector; public SingleConnectionBrokerRequestHandler(PinotConfiguration config, String brokerId, - BrokerRoutingManager routingManager, AccessControlFactory accessControlFactory, + RoutingManager routingManager, AccessControlFactory accessControlFactory, QueryQuotaManager queryQuotaManager, TableCache tableCache, NettyConfig nettyConfig, TlsConfig tlsConfig, ServerRoutingStatsManager serverRoutingStatsManager, FailureDetector failureDetector, ThreadResourceUsageAccountant accountant) { From 1d03cd89e548dadbba7b7e0a674408eebe4b1112 Mon Sep 17 00:00:00 2001 From: Shaurya Chaturvedi Date: Tue, 5 Aug 2025 09:52:57 -0700 Subject: [PATCH 078/167] [timeseries] Enable table name driven broker selection for time series queries in Pinot Controller (#16509) Co-authored-by: Shaurya Chaturvedi --- pinot-controller/pom.xml | 4 +++ .../api/resources/PinotQueryResource.java | 29 +++++++++++++------ .../planner/TimeSeriesQueryEnvironment.java | 27 +++++++++++------ .../tsdb/spi/TimeSeriesLogicalPlanner.java | 24 +++++++++++++++ 4 files changed, 66 insertions(+), 18 deletions(-) diff --git a/pinot-controller/pom.xml b/pinot-controller/pom.xml index 76125445d247..cb8af6fdd6ab 100644 --- a/pinot-controller/pom.xml +++ b/pinot-controller/pom.xml @@ -50,6 +50,10 @@ org.apache.pinot pinot-csv + + org.apache.pinot + pinot-timeseries-planner + com.101tec diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotQueryResource.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotQueryResource.java index 1e4d4f18863e..d15a8fd61a3d 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotQueryResource.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotQueryResource.java @@ -28,6 +28,7 @@ import java.net.HttpURLConnection; import java.net.URL; import java.nio.charset.StandardCharsets; +import java.time.Duration; import java.util.ArrayList; import java.util.HashSet; import java.util.List; @@ -81,6 +82,11 @@ import org.apache.pinot.sql.parsers.CalciteSqlParser; import org.apache.pinot.sql.parsers.PinotSqlType; import org.apache.pinot.sql.parsers.SqlNodeAndOptions; +import org.apache.pinot.tsdb.planner.TimeSeriesQueryEnvironment; +import org.apache.pinot.tsdb.planner.TimeSeriesTableMetadataProvider; +import org.apache.pinot.tsdb.spi.RangeTimeSeriesRequest; +import org.apache.pinot.tsdb.spi.TimeSeriesLogicalPlanResult; +import org.apache.pinot.tsdb.spi.TimeSeriesLogicalPlanner; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -495,17 +501,22 @@ private StreamingOutput executeTimeSeriesQueryCatching(HttpHeaders httpHeaders, } } + private String retrieveBrokerForTimeSeriesQuery(String query, String language, String start, String end) { + TimeSeriesLogicalPlanner planner = TimeSeriesQueryEnvironment.buildLogicalPlanner(language, _controllerConf); + TimeSeriesLogicalPlanResult planResult = planner.plan( + new RangeTimeSeriesRequest(language, query, Integer.parseInt(start), Long.parseLong(end), + 60L, Duration.ofMinutes(1), 100, 100, ""), + new TimeSeriesTableMetadataProvider(_pinotHelixResourceManager.getTableCache())); + String tableName = planner.getTableName(planResult); + String rawTableName = TableNameBuilder.extractRawTableName(tableName); + List instanceIds = _pinotHelixResourceManager.getBrokerInstancesFor(rawTableName); + return selectRandomInstanceId(instanceIds); + } + private StreamingOutput executeTimeSeriesQuery(HttpHeaders httpHeaders, String language, String query, - String start, String end, String step) throws Exception { + String start, String end, String step) throws Exception { LOGGER.debug("Language: {}, Query: {}, Start: {}, End: {}, Step: {}", language, query, start, end, step); - - // Get available broker instances for timeseries queries - List instanceIds = _pinotHelixResourceManager.getAllBrokerInstances(); - if (instanceIds.isEmpty()) { - throw QueryErrorCode.BROKER_INSTANCE_MISSING.asException("No online broker found for timeseries query"); - } - - String instanceId = selectRandomInstanceId(instanceIds); + String instanceId = retrieveBrokerForTimeSeriesQuery(query, language, start, end); return sendTimeSeriesRequestToBroker(language, query, start, end, step, instanceId, httpHeaders); } diff --git a/pinot-timeseries/pinot-timeseries-planner/src/main/java/org/apache/pinot/tsdb/planner/TimeSeriesQueryEnvironment.java b/pinot-timeseries/pinot-timeseries-planner/src/main/java/org/apache/pinot/tsdb/planner/TimeSeriesQueryEnvironment.java index a0246ff57cb8..6cc12ebf0a2c 100644 --- a/pinot-timeseries/pinot-timeseries-planner/src/main/java/org/apache/pinot/tsdb/planner/TimeSeriesQueryEnvironment.java +++ b/pinot-timeseries/pinot-timeseries-planner/src/main/java/org/apache/pinot/tsdb/planner/TimeSeriesQueryEnvironment.java @@ -61,16 +61,8 @@ public void init(PinotConfiguration config) { .split(","); LOGGER.info("Found {} configured time series languages. List: {}", languages.length, languages); for (String language : languages) { - String configPrefix = PinotTimeSeriesConfiguration.getLogicalPlannerConfigKey(language); - String klassName = - config.getProperty(PinotTimeSeriesConfiguration.getLogicalPlannerConfigKey(language)); - Preconditions.checkNotNull(klassName, "Logical planner class not found for language: " + language); - // Create the planner with empty constructor try { - Class klass = TimeSeriesQueryEnvironment.class.getClassLoader().loadClass(klassName); - Constructor constructor = klass.getConstructor(); - TimeSeriesLogicalPlanner planner = (TimeSeriesLogicalPlanner) constructor.newInstance(); - planner.init(config.subset(configPrefix)); + TimeSeriesLogicalPlanner planner = buildLogicalPlanner(language, config); _plannerMap.put(language, planner); } catch (Exception e) { throw new RuntimeException("Failed to instantiate logical planner for language: " + language, e); @@ -80,6 +72,23 @@ public void init(PinotConfiguration config) { TableScanVisitor.INSTANCE.init(_routingManager, new ImplicitHybridTableRouteProvider(), _tableCache); } + public static TimeSeriesLogicalPlanner buildLogicalPlanner(String language, PinotConfiguration config) + throws RuntimeException { + String configPrefix = PinotTimeSeriesConfiguration.getLogicalPlannerConfigKey(language); + String klassName = config.getProperty(configPrefix); + Preconditions.checkNotNull(klassName, "Logical planner class not found for language: " + language); + // Create the planner with empty constructor + try { + Class klass = TimeSeriesQueryEnvironment.class.getClassLoader().loadClass(klassName); + Constructor constructor = klass.getConstructor(); + TimeSeriesLogicalPlanner planner = (TimeSeriesLogicalPlanner) constructor.newInstance(); + planner.init(config.subset(configPrefix)); + return planner; + } catch (Exception e) { + throw new RuntimeException("Failed to instantiate logical planner for language: " + language, e); + } + } + public TimeSeriesLogicalPlanResult buildLogicalPlan(RangeTimeSeriesRequest request) { Preconditions.checkState(_plannerMap.containsKey(request.getLanguage()), "No logical planner found for engine: %s. Available: %s", request.getLanguage(), diff --git a/pinot-timeseries/pinot-timeseries-spi/src/main/java/org/apache/pinot/tsdb/spi/TimeSeriesLogicalPlanner.java b/pinot-timeseries/pinot-timeseries-spi/src/main/java/org/apache/pinot/tsdb/spi/TimeSeriesLogicalPlanner.java index 0a5b9d907a98..e6ab2f8e2615 100644 --- a/pinot-timeseries/pinot-timeseries-spi/src/main/java/org/apache/pinot/tsdb/spi/TimeSeriesLogicalPlanner.java +++ b/pinot-timeseries/pinot-timeseries-spi/src/main/java/org/apache/pinot/tsdb/spi/TimeSeriesLogicalPlanner.java @@ -18,6 +18,7 @@ */ package org.apache.pinot.tsdb.spi; +import java.util.Stack; import org.apache.pinot.spi.env.PinotConfiguration; import org.apache.pinot.tsdb.spi.plan.BaseTimeSeriesPlanNode; import org.apache.pinot.tsdb.spi.plan.LeafTimeSeriesPlanNode; @@ -35,4 +36,27 @@ public interface TimeSeriesLogicalPlanner { void init(PinotConfiguration pinotConfiguration); TimeSeriesLogicalPlanResult plan(RangeTimeSeriesRequest request, TimeSeriesMetadata metadata); + + /** + * Returns the name of the table from the logical plan result by traversing the plan tree and extracting the + * table name from the first encountered {@link LeafTimeSeriesPlanNode} + * This method is recommended to be overriden by implementations for more efficient table name extraction. + */ + default String getTableName(TimeSeriesLogicalPlanResult result) { + BaseTimeSeriesPlanNode node = result.getPlanNode(); + + Stack nodeStack = new Stack<>(); + nodeStack.push(node); + + while (!nodeStack.isEmpty()) { + BaseTimeSeriesPlanNode currentNode = nodeStack.pop(); + if (currentNode instanceof LeafTimeSeriesPlanNode) { + return ((LeafTimeSeriesPlanNode) currentNode).getTableName(); + } + for (BaseTimeSeriesPlanNode child : currentNode.getInputs()) { + nodeStack.push(child); + } + } + throw new RuntimeException("No LeafTimeSeriesPlanNode found in the plan tree."); + } } From e961bdfc55044fdd9625b0932a7aa9a27ec3adaa Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 5 Aug 2025 14:52:32 -0700 Subject: [PATCH 079/167] Bump software.amazon.awssdk:bom from 2.32.14 to 2.32.15 (#16512) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 4c9319a2b0eb..53098131c6e3 100644 --- a/pom.xml +++ b/pom.xml @@ -178,7 +178,7 @@ 0.15.1 0.4.7 4.2.2 - 2.32.14 + 2.32.15 1.2.37 1.22.0 2.14.0 From c9cf6d094babd38d31bae913c6232765c06dd394 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Tue, 5 Aug 2025 14:53:06 -0700 Subject: [PATCH 080/167] Bump immutables.version from 2.11.1 to 2.11.2 (#16513) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 53098131c6e3..a284af85d8be 100644 --- a/pom.xml +++ b/pom.xml @@ -157,7 +157,7 @@ 2.6.0 2.5.0 1.40.0 - 2.11.1 + 2.11.2 9.12.0 0.10.2 0.25.0 From 575032cfd19dc2d37bf38d974a0bbb40acafd987 Mon Sep 17 00:00:00 2001 From: Jhow <44998515+J-HowHuang@users.noreply.github.com> Date: Tue, 5 Aug 2025 14:59:45 -0700 Subject: [PATCH 081/167] [Bug fix] Update logic that determines if a table is associated with a tenant (#16477) --- .../resources/PinotTenantRestletResource.java | 4 +- .../tenant/DefaultTenantRebalancer.java | 4 +- .../segment/local/utils/TableConfigUtils.java | 7 ++ .../local/utils/TableConfigUtilsTest.java | 77 +++++++++++++++++++ 4 files changed, 86 insertions(+), 6 deletions(-) diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTenantRestletResource.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTenantRestletResource.java index fba07869d959..817062d8a19e 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTenantRestletResource.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTenantRestletResource.java @@ -60,7 +60,6 @@ import org.apache.pinot.common.assignment.InstancePartitionsUtils; import org.apache.pinot.common.metrics.ControllerMeter; import org.apache.pinot.common.metrics.ControllerMetrics; -import org.apache.pinot.common.utils.config.TagNameUtils; import org.apache.pinot.controller.api.access.AccessType; import org.apache.pinot.controller.api.access.Authenticate; import org.apache.pinot.controller.api.exception.ControllerApplicationException; @@ -409,8 +408,7 @@ private String getTablesServedFromServerTenant(String tenantName, @Nullable Stri LOGGER.error("Unable to retrieve table config for table: {}", tableNameWithType); continue; } - Set relevantTags = TableConfigUtils.getRelevantTags(tableConfig); - if (relevantTags.contains(TagNameUtils.getServerTagForTenant(tenantName, tableConfig.getTableType()))) { + if (TableConfigUtils.isRelevantToTenant(tableConfig, tenantName)) { tables.add(tableNameWithType); if (withTableProperties) { IdealState idealState = _pinotHelixResourceManager.getTableIdealState(tableNameWithType); diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/DefaultTenantRebalancer.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/DefaultTenantRebalancer.java index a23bf2624166..749afc72e2f2 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/DefaultTenantRebalancer.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/tenant/DefaultTenantRebalancer.java @@ -32,7 +32,6 @@ import javax.annotation.Nullable; import org.apache.commons.lang3.tuple.Pair; import org.apache.pinot.common.exception.TableNotFoundException; -import org.apache.pinot.common.utils.config.TagNameUtils; import org.apache.pinot.controller.helix.core.PinotHelixResourceManager; import org.apache.pinot.controller.helix.core.rebalance.RebalanceConfig; import org.apache.pinot.controller.helix.core.rebalance.RebalanceResult; @@ -186,8 +185,7 @@ Set getTenantTables(String tenantName) { LOGGER.error("Unable to retrieve table config for table: {}", table); continue; } - Set relevantTags = TableConfigUtils.getRelevantTags(tableConfig); - if (relevantTags.contains(TagNameUtils.getServerTagForTenant(tenantName, tableConfig.getTableType()))) { + if (TableConfigUtils.isRelevantToTenant(tableConfig, tenantName)) { tables.add(table); } } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java index fa791b811e84..cd4833dd044d 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/TableConfigUtils.java @@ -1735,4 +1735,11 @@ public static Set getRelevantTags(TableConfig tableConfig) { } return relevantTags; } + + public static boolean isRelevantToTenant(TableConfig tableConfig, String tenantName) { + Set relevantTenants = + getRelevantTags(tableConfig).stream().map(TagNameUtils::getTenantFromTag).collect( + Collectors.toSet()); + return relevantTenants.contains(tenantName); + } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/TableConfigUtilsTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/TableConfigUtilsTest.java index bc6f1a29cc08..ee25dc87666e 100644 --- a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/TableConfigUtilsTest.java +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/utils/TableConfigUtilsTest.java @@ -82,6 +82,7 @@ import org.apache.pinot.spi.utils.JsonUtils; import org.apache.pinot.spi.utils.builder.TableConfigBuilder; import org.mockito.Mockito; +import org.testng.annotations.DataProvider; import org.testng.annotations.Test; import static org.mockito.Mockito.when; @@ -3211,4 +3212,80 @@ public void testGetPartitionColumnWithReplicaGroupConfig() { assertEquals(TableConfigUtils.getPartitionColumn(tableConfig), PARTITION_COLUMN); } + + @DataProvider(name = "tableTypeTestDataProvider") + public Object[][] tableTypeTestDataProvider() { + return new Object[][]{ + {TableType.OFFLINE}, + {TableType.REALTIME} + }; + } + + @Test(dataProvider = "tableTypeTestDataProvider") + public void testIsRelevantToTenant(TableType tableType) { + // Test 1: Basic server tenant relevance + TableConfig tableConfig = new TableConfigBuilder(tableType) + .setTableName(TABLE_NAME) + .setServerTenant("serverTenant") + .setBrokerTenant("brokerTenant") + .build(); + + assertTrue(TableConfigUtils.isRelevantToTenant(tableConfig, "serverTenant")); + assertFalse(TableConfigUtils.isRelevantToTenant(tableConfig, "otherTenant")); + assertFalse(TableConfigUtils.isRelevantToTenant(tableConfig, + "brokerTenant")); // broker tenant not relevant for server operations + + // Test 2: Tag override configs + TagOverrideConfig tagOverrideConfig = new TagOverrideConfig("customConsuming_REALTIME", "customCompleted_OFFLINE"); + tableConfig = new TableConfigBuilder(tableType) + .setTableName(TABLE_NAME) + .setTagOverrideConfig(tagOverrideConfig) + .setServerTenant("serverTenant") + .setBrokerTenant("brokerTenant") + .build(); + + assertTrue(TableConfigUtils.isRelevantToTenant(tableConfig, "serverTenant")); + assertTrue(TableConfigUtils.isRelevantToTenant(tableConfig, "customConsuming")); + assertTrue(TableConfigUtils.isRelevantToTenant(tableConfig, "customCompleted")); + assertFalse(TableConfigUtils.isRelevantToTenant(tableConfig, "otherTenant")); + + // Test 3: Instance assignment configs + String tagSuffix = tableType == TableType.OFFLINE ? "_OFFLINE" : "_REALTIME"; + InstanceTagPoolConfig tagPoolConfig = new InstanceTagPoolConfig("tag" + tagSuffix, false, 0, null); + InstanceReplicaGroupPartitionConfig replicaGroupPartitionConfig = + new InstanceReplicaGroupPartitionConfig(false, 0, 0, 0, 0, 0, false, null); + InstanceAssignmentConfig instanceAssignmentConfig = + new InstanceAssignmentConfig(tagPoolConfig, null, replicaGroupPartitionConfig, null, false); + + Map instanceAssignmentConfigMap = new HashMap<>(); + instanceAssignmentConfigMap.put(tableType.name(), instanceAssignmentConfig); + + tableConfig = new TableConfigBuilder(tableType) + .setTableName(TABLE_NAME) + .setServerTenant("serverTenant") + .setBrokerTenant("brokerTenant") + .setInstanceAssignmentConfigMap(instanceAssignmentConfigMap) + .build(); + + assertTrue(TableConfigUtils.isRelevantToTenant(tableConfig, "serverTenant")); + assertTrue(TableConfigUtils.isRelevantToTenant(tableConfig, "tag")); + assertFalse(TableConfigUtils.isRelevantToTenant(tableConfig, "otherTenant")); + + // Test 4: Tier configs + TierConfig tierConfig = + new TierConfig("tier", TierFactory.TIME_SEGMENT_SELECTOR_TYPE.toLowerCase(), "40d", null, + TierFactory.PINOT_SERVER_STORAGE_TYPE, "tierTag" + tagSuffix, null, null); + List tierConfigs = Collections.singletonList(tierConfig); + + tableConfig = new TableConfigBuilder(tableType) + .setTableName(TABLE_NAME) + .setServerTenant("serverTenant") + .setBrokerTenant("brokerTenant") + .setTierConfigList(tierConfigs) + .build(); + + assertTrue(TableConfigUtils.isRelevantToTenant(tableConfig, "serverTenant")); + assertTrue(TableConfigUtils.isRelevantToTenant(tableConfig, "tierTag")); + assertFalse(TableConfigUtils.isRelevantToTenant(tableConfig, "otherTenant")); + } } From 850b6ac28214e5493bbe004409d49675402d74f1 Mon Sep 17 00:00:00 2001 From: Alberto Bastos Date: Wed, 6 Aug 2025 00:09:27 +0200 Subject: [PATCH 082/167] Add prune option for lineage operations (#16310) --- .../common/utils/FileUploadDownloadClient.java | 6 ++++-- .../PinotSegmentUploadDownloadRestletResource.java | 6 ++++++ .../helix/core/retention/RetentionManager.java | 3 ++- .../minion/tasks/SegmentConversionUtils.java | 14 +++++++++++--- .../local/utils/ConsistentDataPushUtils.java | 2 +- 5 files changed, 24 insertions(+), 7 deletions(-) diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/FileUploadDownloadClient.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/FileUploadDownloadClient.java index 543f6c80f74d..1912d75666af 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/FileUploadDownloadClient.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/FileUploadDownloadClient.java @@ -126,6 +126,7 @@ public static FileUploadType getDefaultUploadType() { private static final String SEGMENT_LINEAGE_ENTRY_ID_PARAMETER = "&segmentLineageEntryId="; private static final String FORCE_REVERT_PARAMETER = "&forceRevert="; private static final String FORCE_CLEANUP_PARAMETER = "&forceCleanup="; + private static final String CLEANUP_PARAMETER = "&cleanup="; private static final String RETENTION_PARAMETER = "retention="; @@ -391,11 +392,12 @@ public static URI getStartReplaceSegmentsURI(URI controllerURI, String rawTableN } public static URI getEndReplaceSegmentsURI(URI controllerURI, String rawTableName, String tableType, - String segmentLineageEntryId) + String segmentLineageEntryId, boolean cleanup) throws URISyntaxException { return getURI(controllerURI.getScheme(), controllerURI.getHost(), controllerURI.getPort(), OLD_SEGMENT_PATH + "/" + rawTableName + END_REPLACE_SEGMENTS_PATH, - TYPE_DELIMITER + tableType + SEGMENT_LINEAGE_ENTRY_ID_PARAMETER + segmentLineageEntryId); + TYPE_DELIMITER + tableType + + SEGMENT_LINEAGE_ENTRY_ID_PARAMETER + segmentLineageEntryId + CLEANUP_PARAMETER + cleanup); } public static URI getRevertReplaceSegmentsURI(URI controllerURI, String rawTableName, String tableType, diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java index 4cb0fcfde7f7..d4e07f364aab 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentUploadDownloadRestletResource.java @@ -95,6 +95,7 @@ import org.apache.pinot.controller.api.upload.SegmentValidationUtils; import org.apache.pinot.controller.api.upload.ZKOperator; import org.apache.pinot.controller.helix.core.PinotHelixResourceManager; +import org.apache.pinot.controller.helix.core.retention.RetentionManager; import org.apache.pinot.controller.validation.StorageQuotaChecker; import org.apache.pinot.core.auth.Actions; import org.apache.pinot.core.auth.Authorize; @@ -1019,6 +1020,8 @@ public Response endReplaceSegments( @ApiParam(value = "OFFLINE|REALTIME", required = true) @QueryParam("type") String tableTypeStr, @ApiParam(value = "Segment lineage entry id returned by startReplaceSegments API", required = true) @QueryParam("segmentLineageEntryId") String segmentLineageEntryId, + @ApiParam(value = "Trigger an immediate segment cleanup") @QueryParam("cleanup") @DefaultValue("false") + boolean cleanupSegments, @ApiParam(value = "Fields belonging to end replace segment request") EndReplaceSegmentsRequest endReplaceSegmentsRequest, @Context HttpHeaders headers) { tableName = DatabaseUtils.translateTableName(tableName, headers); @@ -1034,6 +1037,9 @@ public Response endReplaceSegments( Preconditions.checkNotNull(segmentLineageEntryId, "'segmentLineageEntryId' should not be null"); _pinotHelixResourceManager.endReplaceSegments(tableNameWithType, segmentLineageEntryId, endReplaceSegmentsRequest); + if (cleanupSegments) { + _pinotHelixResourceManager.invokeControllerPeriodicTask(tableNameWithType, RetentionManager.TASK_NAME, null); + } return Response.ok().build(); } catch (Exception e) { _controllerMetrics.addMeteredTableValue(tableNameWithType, ControllerMeter.NUMBER_END_REPLACE_FAILURE, 1); diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/retention/RetentionManager.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/retention/RetentionManager.java index d647574eee41..3f6c162b81a3 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/retention/RetentionManager.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/retention/RetentionManager.java @@ -75,6 +75,7 @@ *

    It is scheduled to run only on leader controller. */ public class RetentionManager extends ControllerPeriodicTask { + public static final String TASK_NAME = "RetentionManager"; public static final long OLD_LLC_SEGMENTS_RETENTION_IN_MILLIS = TimeUnit.DAYS.toMillis(5L); public static final int DEFAULT_UNTRACKED_SEGMENTS_DELETION_BATCH_SIZE = 100; private static final RetryPolicy DEFAULT_RETRY_POLICY = RetryPolicies.randomDelayRetryPolicy(20, 100L, 200L); @@ -87,7 +88,7 @@ public class RetentionManager extends ControllerPeriodicTask { public RetentionManager(PinotHelixResourceManager pinotHelixResourceManager, LeadControllerManager leadControllerManager, ControllerConf config, ControllerMetrics controllerMetrics, BrokerServiceHelper brokerServiceHelper) { - super("RetentionManager", config.getRetentionControllerFrequencyInSeconds(), + super(TASK_NAME, config.getRetentionControllerFrequencyInSeconds(), config.getRetentionManagerInitialDelayInSeconds(), pinotHelixResourceManager, leadControllerManager, controllerMetrics); _untrackedSegmentDeletionEnabled = config.getUntrackedSegmentDeletionEnabled(); diff --git a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/SegmentConversionUtils.java b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/SegmentConversionUtils.java index c0dad3353874..7f5f2b0f03d9 100644 --- a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/SegmentConversionUtils.java +++ b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/SegmentConversionUtils.java @@ -204,19 +204,27 @@ public static String startSegmentReplace(String tableNameWithType, String upload public static void endSegmentReplace(String tableNameWithType, String uploadURL, String segmentLineageEntryId, int socketTimeoutMs, @Nullable AuthProvider authProvider) throws Exception { - endSegmentReplace(tableNameWithType, uploadURL, null, segmentLineageEntryId, socketTimeoutMs, authProvider); + endSegmentReplace(tableNameWithType, uploadURL, null, segmentLineageEntryId, + socketTimeoutMs, authProvider, false); + } + + public static void endSegmentReplace(String tableNameWithType, String uploadURL, String segmentLineageEntryId, + int socketTimeoutMs, @Nullable AuthProvider authProvider, boolean cleanup) + throws Exception { + endSegmentReplace(tableNameWithType, uploadURL, null, segmentLineageEntryId, + socketTimeoutMs, authProvider, cleanup); } public static void endSegmentReplace(String tableNameWithType, String uploadURL, @Nullable EndReplaceSegmentsRequest endReplaceSegmentsRequest, String segmentLineageEntryId, int socketTimeoutMs, - @Nullable AuthProvider authProvider) + @Nullable AuthProvider authProvider, boolean cleanup) throws Exception { String rawTableName = TableNameBuilder.extractRawTableName(tableNameWithType); TableType tableType = TableNameBuilder.getTableTypeFromTableName(tableNameWithType); SSLContext sslContext = MinionContext.getInstance().getSSLContext(); try (FileUploadDownloadClient fileUploadDownloadClient = new FileUploadDownloadClient(sslContext)) { URI uri = FileUploadDownloadClient.getEndReplaceSegmentsURI(new URI(uploadURL), rawTableName, tableType.name(), - segmentLineageEntryId); + segmentLineageEntryId, cleanup); SimpleHttpResponse response = fileUploadDownloadClient.endReplaceSegments(uri, socketTimeoutMs, endReplaceSegmentsRequest, authProvider); LOGGER.info("Got response {}: {} while sending end replace segment request for table: {}, uploadURL: {}, request:" diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/ConsistentDataPushUtils.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/ConsistentDataPushUtils.java index f7a07083fc87..0bdbf71f6175 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/ConsistentDataPushUtils.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/utils/ConsistentDataPushUtils.java @@ -164,7 +164,7 @@ public static void endReplaceSegments(SegmentGenerationJobSpec spec, Map { try { SimpleHttpResponse response = From 64a07ec3ee9ff6825605e6a07461fa30cb5e58dc Mon Sep 17 00:00:00 2001 From: Song Fu <131259315+songwdfu@users.noreply.github.com> Date: Tue, 5 Aug 2025 16:33:04 -0700 Subject: [PATCH 083/167] [Planner] Add ProjectRemove rule into PRUNE_RULES (#16518) --- .../calcite/rel/rules/PinotQueryRuleSets.java | 2 + .../src/test/resources/queries/JoinPlans.json | 65 +++++++++---------- 2 files changed, 32 insertions(+), 35 deletions(-) diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java index 191a20058845..e07de5a53e9a 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java @@ -183,6 +183,8 @@ private PinotQueryRuleSets() { .withDescription(PlannerRuleNames.AGGREGATE_PROJECT_MERGE).toRule(), ProjectMergeRule.Config.DEFAULT .withDescription(PlannerRuleNames.PROJECT_MERGE).toRule(), + ProjectRemoveRule.Config.DEFAULT + .withDescription(PlannerRuleNames.PROJECT_REMOVE).toRule(), FilterMergeRule.Config.DEFAULT .withDescription(PlannerRuleNames.FILTER_MERGE).toRule(), AggregateRemoveRule.Config.DEFAULT diff --git a/pinot-query-planner/src/test/resources/queries/JoinPlans.json b/pinot-query-planner/src/test/resources/queries/JoinPlans.json index dea372f5fdc6..c1fb4b1e52bd 100644 --- a/pinot-query-planner/src/test/resources/queries/JoinPlans.json +++ b/pinot-query-planner/src/test/resources/queries/JoinPlans.json @@ -379,21 +379,19 @@ "\n LogicalFilter(condition=[=($0, _UTF-8'foo')])", "\n PinotLogicalTableScan(table=[[default, b]])", "\n PinotLogicalExchange(distribution=[hash[0]])", - "\n LogicalProject(col3=[$0], $f1=[$1])", - "\n PinotLogicalAggregate(group=[{0}], agg#0=[MIN($1)], aggType=[FINAL])", - "\n PinotLogicalExchange(distribution=[hash[0]])", - "\n PinotLogicalAggregate(group=[{0}], agg#0=[MIN($1)], aggType=[LEAF])", - "\n LogicalProject(col3=[$2], $f1=[true])", - "\n LogicalFilter(condition=[=($0, _UTF-8'bar')])", - "\n PinotLogicalTableScan(table=[[default, b]])", + "\n PinotLogicalAggregate(group=[{0}], agg#0=[MIN($1)], aggType=[FINAL])", + "\n PinotLogicalExchange(distribution=[hash[0]])", + "\n PinotLogicalAggregate(group=[{0}], agg#0=[MIN($1)], aggType=[LEAF])", + "\n LogicalProject(col3=[$2], $f1=[true])", + "\n LogicalFilter(condition=[=($0, _UTF-8'bar')])", + "\n PinotLogicalTableScan(table=[[default, b]])", "\n PinotLogicalExchange(distribution=[hash[0]])", - "\n LogicalProject(col3=[$0], $f1=[$1])", - "\n PinotLogicalAggregate(group=[{0}], agg#0=[MIN($1)], aggType=[FINAL])", - "\n PinotLogicalExchange(distribution=[hash[0]])", - "\n PinotLogicalAggregate(group=[{0}], agg#0=[MIN($1)], aggType=[LEAF])", - "\n LogicalProject(col3=[$2], $f1=[true])", - "\n LogicalFilter(condition=[=($0, _UTF-8'foobar')])", - "\n PinotLogicalTableScan(table=[[default, b]])", + "\n PinotLogicalAggregate(group=[{0}], agg#0=[MIN($1)], aggType=[FINAL])", + "\n PinotLogicalExchange(distribution=[hash[0]])", + "\n PinotLogicalAggregate(group=[{0}], agg#0=[MIN($1)], aggType=[LEAF])", + "\n LogicalProject(col3=[$2], $f1=[true])", + "\n LogicalFilter(condition=[=($0, _UTF-8'foobar')])", + "\n PinotLogicalTableScan(table=[[default, b]])", "\n" ] }, @@ -634,29 +632,26 @@ "\n LogicalFilter(condition=[AND(=($0, _UTF-8'foo'), =($1, _UTF-8'xylo'), =($3, 12.0000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000), NOT($4))])", "\n PinotLogicalTableScan(table=[[default, a]])", "\n PinotLogicalExchange(distribution=[hash[0]])", - "\n LogicalProject(col3=[$0], $f1=[$1])", - "\n PinotLogicalAggregate(group=[{0}], agg#0=[MIN($1)], aggType=[FINAL])", - "\n PinotLogicalExchange(distribution=[hash[0]])", - "\n PinotLogicalAggregate(group=[{0}], agg#0=[MIN($1)], aggType=[LEAF])", - "\n LogicalProject(col3=[$2], $f1=[true])", - "\n LogicalFilter(condition=[=($0, _UTF-8'foo')])", - "\n PinotLogicalTableScan(table=[[default, b]])", + "\n PinotLogicalAggregate(group=[{0}], agg#0=[MIN($1)], aggType=[FINAL])", + "\n PinotLogicalExchange(distribution=[hash[0]])", + "\n PinotLogicalAggregate(group=[{0}], agg#0=[MIN($1)], aggType=[LEAF])", + "\n LogicalProject(col3=[$2], $f1=[true])", + "\n LogicalFilter(condition=[=($0, _UTF-8'foo')])", + "\n PinotLogicalTableScan(table=[[default, b]])", "\n PinotLogicalExchange(distribution=[hash[0]])", - "\n LogicalProject(col3=[$0], $f1=[$1])", - "\n PinotLogicalAggregate(group=[{0}], agg#0=[MIN($1)], aggType=[FINAL])", - "\n PinotLogicalExchange(distribution=[hash[0]])", - "\n PinotLogicalAggregate(group=[{0}], agg#0=[MIN($1)], aggType=[LEAF])", - "\n LogicalProject(col3=[$2], $f1=[true])", - "\n LogicalFilter(condition=[=($0, _UTF-8'bar')])", - "\n PinotLogicalTableScan(table=[[default, b]])", + "\n PinotLogicalAggregate(group=[{0}], agg#0=[MIN($1)], aggType=[FINAL])", + "\n PinotLogicalExchange(distribution=[hash[0]])", + "\n PinotLogicalAggregate(group=[{0}], agg#0=[MIN($1)], aggType=[LEAF])", + "\n LogicalProject(col3=[$2], $f1=[true])", + "\n LogicalFilter(condition=[=($0, _UTF-8'bar')])", + "\n PinotLogicalTableScan(table=[[default, b]])", "\n PinotLogicalExchange(distribution=[hash[0]])", - "\n LogicalProject(col3=[$0], $f1=[$1])", - "\n PinotLogicalAggregate(group=[{0}], agg#0=[MIN($1)], aggType=[FINAL])", - "\n PinotLogicalExchange(distribution=[hash[0]])", - "\n PinotLogicalAggregate(group=[{0}], agg#0=[MIN($1)], aggType=[LEAF])", - "\n LogicalProject(col3=[$2], $f1=[true])", - "\n LogicalFilter(condition=[=($0, _UTF-8'foobar')])", - "\n PinotLogicalTableScan(table=[[default, b]])", + "\n PinotLogicalAggregate(group=[{0}], agg#0=[MIN($1)], aggType=[FINAL])", + "\n PinotLogicalExchange(distribution=[hash[0]])", + "\n PinotLogicalAggregate(group=[{0}], agg#0=[MIN($1)], aggType=[LEAF])", + "\n LogicalProject(col3=[$2], $f1=[true])", + "\n LogicalFilter(condition=[=($0, _UTF-8'foobar')])", + "\n PinotLogicalTableScan(table=[[default, b]])", "\n PinotLogicalExchange(distribution=[hash[0]])", "\n LogicalProject(col3=[$2])", "\n LogicalFilter(condition=[=($0, _UTF-8'fork')])", From 407d6f24ab7e81a81b8f31fffa87a39204c10368 Mon Sep 17 00:00:00 2001 From: Song Fu <131259315+songwdfu@users.noreply.github.com> Date: Tue, 5 Aug 2025 23:38:33 -0700 Subject: [PATCH 084/167] Clear threadLocal maps after use for DictionaryBasedGroupKeyGenerator (#16454) --- .../combine/GroupByCombineOperator.java | 27 ++++---- .../query/FilteredGroupByOperator.java | 2 + .../core/operator/query/GroupByOperator.java | 2 + .../groupby/AggregationGroupByResult.java | 7 ++ .../DictionaryBasedGroupKeyGenerator.java | 64 +++++++++++++------ .../groupby/GroupKeyGenerator.java | 7 +- .../DictionaryBasedGroupKeyGeneratorTest.java | 25 ++++++++ 7 files changed, 102 insertions(+), 32 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/GroupByCombineOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/GroupByCombineOperator.java index 18e67b989976..0f39e7b8815d 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/GroupByCombineOperator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/combine/GroupByCombineOperator.java @@ -143,18 +143,23 @@ protected void processSegments() { AggregationGroupByResult aggregationGroupByResult = resultsBlock.getAggregationGroupByResult(); if (aggregationGroupByResult != null) { // Iterate over the group-by keys, for each key, update the group-by result in the indexedTable - Iterator dicGroupKeyIterator = aggregationGroupByResult.getGroupKeyIterator(); - while (dicGroupKeyIterator.hasNext()) { - GroupKeyGenerator.GroupKey groupKey = dicGroupKeyIterator.next(); - Object[] keys = groupKey._keys; - Object[] values = Arrays.copyOf(keys, _numColumns); - int groupId = groupKey._groupId; - for (int i = 0; i < _numAggregationFunctions; i++) { - values[_numGroupByExpressions + i] = aggregationGroupByResult.getResultForGroupId(i, groupId); + try { + Iterator dicGroupKeyIterator = aggregationGroupByResult.getGroupKeyIterator(); + while (dicGroupKeyIterator.hasNext()) { + GroupKeyGenerator.GroupKey groupKey = dicGroupKeyIterator.next(); + Object[] keys = groupKey._keys; + Object[] values = Arrays.copyOf(keys, _numColumns); + int groupId = groupKey._groupId; + for (int i = 0; i < _numAggregationFunctions; i++) { + values[_numGroupByExpressions + i] = aggregationGroupByResult.getResultForGroupId(i, groupId); + } + _indexedTable.upsert(new Key(keys), new Record(values)); + Tracing.ThreadAccountantOps.sampleAndCheckInterruptionPeriodically(mergedKeys); + mergedKeys++; } - _indexedTable.upsert(new Key(keys), new Record(values)); - Tracing.ThreadAccountantOps.sampleAndCheckInterruptionPeriodically(mergedKeys); - mergedKeys++; + } finally { + // Release the resources used by the group key generator + aggregationGroupByResult.closeGroupKeyGenerator(); } } } else { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/FilteredGroupByOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/FilteredGroupByOperator.java index 1afda13ca68e..af438a5eb30d 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/FilteredGroupByOperator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/FilteredGroupByOperator.java @@ -195,6 +195,8 @@ protected GroupByResultsBlock getNextBlock() { TableResizer tableResizer = new TableResizer(_dataSchema, _queryContext); Collection intermediateRecords = tableResizer.trimInSegmentResults(groupKeyGenerator, groupByResultHolders, trimSize); + // Release the resources used by the group key generator + groupKeyGenerator.close(); ServerMetrics.get().addMeteredGlobalValue(ServerMeter.AGGREGATE_TIMES_GROUPS_TRIMMED, 1); boolean unsafeTrim = _queryContext.isUnsafeTrim(); // set trim flag only if it's not safe diff --git a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java index 5b704ceb476e..1a15f58e71f1 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java @@ -155,6 +155,8 @@ protected GroupByResultsBlock getNextBlock() { if (groupByExecutor.getNumGroups() > trimSize) { TableResizer tableResizer = new TableResizer(_dataSchema, _queryContext); Collection intermediateRecords = groupByExecutor.trimGroupByResult(trimSize, tableResizer); + // trim groupKeyGenerator after getting intermediateRecords + groupByExecutor.getGroupKeyGenerator().close(); ServerMetrics.get().addMeteredGlobalValue(ServerMeter.AGGREGATE_TIMES_GROUPS_TRIMMED, 1); boolean unsafeTrim = _queryContext.isUnsafeTrim(); // set trim flag only if it's not safe diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/AggregationGroupByResult.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/AggregationGroupByResult.java index 49c2361c3c7a..9f9f3641844d 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/AggregationGroupByResult.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/AggregationGroupByResult.java @@ -51,6 +51,13 @@ public Iterator getGroupKeyIterator() { return _groupKeyGenerator.getGroupKeys(); } + /** + * Clear and trim DictionaryBasedGroupKeyGenerator after use + */ + public void closeGroupKeyGenerator() { + _groupKeyGenerator.close(); + } + public Object getResultForGroupId(int index, int groupId) { return _aggregationFunctions[index].extractGroupByResult(_resultHolders[index], groupId); } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DictionaryBasedGroupKeyGenerator.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DictionaryBasedGroupKeyGenerator.java index 8c55582cb8ba..7fbc0bdb266f 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DictionaryBasedGroupKeyGenerator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DictionaryBasedGroupKeyGenerator.java @@ -140,34 +140,23 @@ public DictionaryBasedGroupKeyGenerator(BaseProjectOperator projectOperator, _isSingleValueColumn[i] = columnContext.isSingleValue(); } if (groupByExpressionSizesFromPredicates != null) { - Pair optimizedCardinality = getOptimizedGroupByCardinality(groupByExpressionSizesFromPredicates, - cardinalityMap); - if (optimizedCardinality.getLeft() && optimizedCardinality.getRight() != null) { - longOverflow = false; - cardinalityProduct = Math.min(optimizedCardinality.getRight(), cardinalityProduct); - } - } - // TODO: Clear the holder after processing the query instead of before + Pair optimizedCardinality = getOptimizedGroupByCardinality(groupByExpressionSizesFromPredicates, + cardinalityMap); + if (optimizedCardinality.getLeft() && optimizedCardinality.getRight() != null) { + longOverflow = false; + cardinalityProduct = Math.min(optimizedCardinality.getRight(), cardinalityProduct); + } + } if (longOverflow) { // ArrayMapBasedHolder _globalGroupIdUpperBound = numGroupsLimit; Object2IntOpenHashMap groupIdMap = THREAD_LOCAL_INT_ARRAY_MAP.get(); - int size = groupIdMap.size(); - groupIdMap.clear(); - if (size > MAX_CACHING_MAP_SIZE) { - groupIdMap.trim(); - } _rawKeyHolder = new ArrayMapBasedHolder(groupIdMap); } else { if (cardinalityProduct > Integer.MAX_VALUE) { // LongMapBasedHolder _globalGroupIdUpperBound = numGroupsLimit; Long2IntOpenHashMap groupIdMap = THREAD_LOCAL_LONG_MAP.get(); - int size = groupIdMap.size(); - groupIdMap.clear(); - if (size > MAX_CACHING_MAP_SIZE) { - groupIdMap.trim(); - } _rawKeyHolder = new LongMapBasedHolder(groupIdMap); } else { _globalGroupIdUpperBound = Math.min((int) cardinalityProduct, numGroupsLimit); @@ -176,7 +165,6 @@ public DictionaryBasedGroupKeyGenerator(BaseProjectOperator projectOperator, if (cardinalityProduct > arrayBasedThreshold || numGroupsLimit < cardinalityProduct) { // IntMapBasedHolder IntGroupIdMap groupIdMap = THREAD_LOCAL_INT_MAP.get(); - groupIdMap.clearAndTrim(); _rawKeyHolder = new IntMapBasedHolder(groupIdMap); } else { _rawKeyHolder = new ArrayBasedHolder(); @@ -245,7 +233,13 @@ public int getNumKeys() { return _rawKeyHolder.getNumKeys(); } - private interface RawKeyHolder { + /// Clear and trim thread-local map of _rawKeyHolder + @Override + public void close() { + _rawKeyHolder.close(); + } + + private interface RawKeyHolder extends AutoCloseable { /** * Process a block of documents for all single-valued group-by columns case. @@ -279,6 +273,9 @@ private interface RawKeyHolder { * Returns current number of unique keys */ int getNumKeys(); + + @Override + void close(); } // This holder works only if it can fit all results, otherwise it fails on AIOOBE or produces too many group keys @@ -286,6 +283,10 @@ private class ArrayBasedHolder implements RawKeyHolder { private final boolean[] _flags = new boolean[_globalGroupIdUpperBound]; private int _numKeys = 0; + @Override + public void close() { + } + @Override public void processSingleValue(int numDocs, int[] outGroupIds) { switch (_numGroupByExpressions) { @@ -416,6 +417,11 @@ public int getNumKeys() { private class IntMapBasedHolder implements RawKeyHolder { private final IntGroupIdMap _groupIdMap; + @Override + public void close() { + _groupIdMap.clearAndTrim(); + } + public IntMapBasedHolder(IntGroupIdMap groupIdMap) { _groupIdMap = groupIdMap; } @@ -633,6 +639,15 @@ public LongMapBasedHolder(Long2IntOpenHashMap groupIdMap) { _groupIdMap = groupIdMap; } + @Override + public void close() { + int size = _groupIdMap.size(); + _groupIdMap.clear(); + if (size > MAX_CACHING_MAP_SIZE) { + _groupIdMap.trim(); + } + } + @Override public void processSingleValue(int numDocs, int[] outGroupIds) { for (int i = 0; i < numDocs; i++) { @@ -813,6 +828,15 @@ public ArrayMapBasedHolder(Object2IntOpenHashMap groupIdMap) { _groupIdMap = groupIdMap; } + @Override + public void close() { + int size = _groupIdMap.size(); + _groupIdMap.clear(); + if (size > MAX_CACHING_MAP_SIZE) { + _groupIdMap.trim(); + } + } + @Override public void processSingleValue(int numDocs, int[] outGroupIds) { for (int i = 0; i < numDocs; i++) { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupKeyGenerator.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupKeyGenerator.java index 393610b4a38a..6d7c5e0edfed 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupKeyGenerator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupKeyGenerator.java @@ -24,8 +24,9 @@ /** * Interface for generating group keys. + * It extends AutoCloseable for thread-local maps to be cleared */ -public interface GroupKeyGenerator { +public interface GroupKeyGenerator extends AutoCloseable { char DELIMITER = '\0'; int INVALID_ID = -1; @@ -75,6 +76,10 @@ public interface GroupKeyGenerator { */ int getNumKeys(); + @Override + default void close() { + } + /** * This class encapsulates the integer group id and the group keys. */ diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/DictionaryBasedGroupKeyGeneratorTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/DictionaryBasedGroupKeyGeneratorTest.java index 8150f968872f..2fbcbc5ee54b 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/DictionaryBasedGroupKeyGeneratorTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/DictionaryBasedGroupKeyGeneratorTest.java @@ -194,6 +194,10 @@ public void testIntMapBasedSingleValue() { assertEquals(dictionaryBasedGroupKeyGenerator.getCurrentGroupKeyUpperBound(), 2, _errorMessage); compareSingleValueBuffer(); testGetGroupKeys(dictionaryBasedGroupKeyGenerator.getGroupKeys(), 2); + + // Test clear and trim + dictionaryBasedGroupKeyGenerator.close(); + assertEquals(DictionaryBasedGroupKeyGenerator.THREAD_LOCAL_INT_MAP.get().size(), 0); } @Test @@ -215,6 +219,10 @@ public void testLongMapBasedSingleValue() { assertEquals(dictionaryBasedGroupKeyGenerator.getCurrentGroupKeyUpperBound(), 2, _errorMessage); compareSingleValueBuffer(); testGetGroupKeys(dictionaryBasedGroupKeyGenerator.getGroupKeys(), 2); + + // Test clear and trim + dictionaryBasedGroupKeyGenerator.close(); + assertEquals(DictionaryBasedGroupKeyGenerator.THREAD_LOCAL_LONG_MAP.get().size(), 0); } @Test @@ -236,6 +244,10 @@ public void testArrayMapBasedSingleValue() { assertEquals(dictionaryBasedGroupKeyGenerator.getCurrentGroupKeyUpperBound(), 2, _errorMessage); compareSingleValueBuffer(); testGetGroupKeys(dictionaryBasedGroupKeyGenerator.getGroupKeys(), 2); + + // Test clear and trim + dictionaryBasedGroupKeyGenerator.close(); + assertEquals(DictionaryBasedGroupKeyGenerator.THREAD_LOCAL_INT_ARRAY_MAP.get().size(), 0); } /** @@ -293,6 +305,10 @@ public void tesIntMapBasedMultiValue() { assertEquals(dictionaryBasedGroupKeyGenerator.getCurrentGroupKeyUpperBound(), numUniqueKeys, _errorMessage); compareMultiValueBuffer(); testGetGroupKeys(dictionaryBasedGroupKeyGenerator.getGroupKeys(), numUniqueKeys); + + // Test clear and trim + dictionaryBasedGroupKeyGenerator.close(); + assertEquals(DictionaryBasedGroupKeyGenerator.THREAD_LOCAL_INT_MAP.get().size(), 0); } @Test @@ -316,6 +332,10 @@ public void testLongMapBasedMultiValue() { assertEquals(dictionaryBasedGroupKeyGenerator.getCurrentGroupKeyUpperBound(), numUniqueKeys, _errorMessage); compareMultiValueBuffer(); testGetGroupKeys(dictionaryBasedGroupKeyGenerator.getGroupKeys(), numUniqueKeys); + + // Test clear and trim + dictionaryBasedGroupKeyGenerator.close(); + assertEquals(DictionaryBasedGroupKeyGenerator.THREAD_LOCAL_LONG_MAP.get().size(), 0); } @Test @@ -338,6 +358,10 @@ public void testArrayMapBasedMultiValue() { assertEquals(dictionaryBasedGroupKeyGenerator.getCurrentGroupKeyUpperBound(), numUniqueKeys, _errorMessage); compareMultiValueBuffer(); testGetGroupKeys(dictionaryBasedGroupKeyGenerator.getGroupKeys(), numUniqueKeys); + + // Test clear and trim + dictionaryBasedGroupKeyGenerator.close(); + assertEquals(DictionaryBasedGroupKeyGenerator.THREAD_LOCAL_INT_ARRAY_MAP.get().size(), 0); } @Test @@ -372,6 +396,7 @@ public void testNumGroupsLimit() { assertEquals(MV_GROUP_KEY_BUFFER[i + 1], MV_GROUP_KEY_BUFFER[1], _errorMessage); } testGetGroupKeys(dictionaryBasedGroupKeyGenerator.getGroupKeys(), numGroupsLimit); + dictionaryBasedGroupKeyGenerator.close(); } private static ExpressionContext[] getExpressions(String[] columns) { From 77d6de9d42ce7455996296f61b4258a4d02415c0 Mon Sep 17 00:00:00 2001 From: "Xiaotian (Jackie) Jiang" <17555551+Jackie-Jiang@users.noreply.github.com> Date: Wed, 6 Aug 2025 01:46:03 -0600 Subject: [PATCH 085/167] Fix string parsing for value aggregator (#16519) --- .../core/data/manager/TableIndexingTest.java | 3 +- .../core/startree/v2/MaxStarTreeV2Test.java | 6 ++-- .../core/startree/v2/MinStarTreeV2Test.java | 6 ++-- .../v2/SumPrecisionStarTreeV2Test.java | 2 +- .../core/startree/v2/SumStarTreeV2Test.java | 6 ++-- .../src/test/resources/TableIndexingTest.csv | 4 +-- .../local/aggregator/AvgValueAggregator.java | 4 +-- .../local/aggregator/MaxValueAggregator.java | 10 +++--- .../MinMaxRangeValueAggregator.java | 5 ++- .../local/aggregator/MinValueAggregator.java | 10 +++--- .../PercentileEstValueAggregator.java | 16 +++++++-- .../PercentileTDigestValueAggregator.java | 4 +-- .../local/aggregator/SumValueAggregator.java | 10 +++--- .../aggregator/ValueAggregatorUtils.java | 36 +++++++++++++++++++ 14 files changed, 85 insertions(+), 37 deletions(-) create mode 100644 pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/ValueAggregatorUtils.java diff --git a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/TableIndexingTest.java b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/TableIndexingTest.java index def39a8bb657..9ca4a67544ab 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/TableIndexingTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/TableIndexingTest.java @@ -523,7 +523,8 @@ public void testAddIndex(TestCase testCase) { Assert.fail("No expected status found for test case: " + testCase); } else if (testCase._expectedSuccess && testCase._error != null) { Assert.fail("Expected success for test case: " + testCase + " but got error: " + testCase._error); - } else if (!testCase._expectedSuccess && !testCase.getErrorMessage().equals(testCase._expectedMessage)) { + } else if (!testCase._expectedSuccess && !testCase.getErrorMessage().equals(testCase._expectedMessage) + && !testCase.getErrorMessage().matches(testCase._expectedMessage)) { Assert.fail("Expected error: \"" + testCase._expectedMessage + "\" for test case: " + testCase + " but got: \"" + testCase.getErrorMessage() + "\""); } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/MaxStarTreeV2Test.java b/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/MaxStarTreeV2Test.java index 1dfcc2cedbfa..fa910a2c3ea5 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/MaxStarTreeV2Test.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/MaxStarTreeV2Test.java @@ -26,10 +26,10 @@ import static org.testng.Assert.assertEquals; -public class MaxStarTreeV2Test extends BaseStarTreeV2Test { +public class MaxStarTreeV2Test extends BaseStarTreeV2Test { @Override - ValueAggregator getValueAggregator() { + ValueAggregator getValueAggregator() { return new MaxValueAggregator(); } @@ -39,7 +39,7 @@ DataType getRawValueType() { } @Override - Number getRandomRawValue(Random random) { + Object getRandomRawValue(Random random) { return random.nextDouble(); } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/MinStarTreeV2Test.java b/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/MinStarTreeV2Test.java index b261c76fa257..8d9bc9c9477a 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/MinStarTreeV2Test.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/MinStarTreeV2Test.java @@ -26,10 +26,10 @@ import static org.testng.Assert.assertEquals; -public class MinStarTreeV2Test extends BaseStarTreeV2Test { +public class MinStarTreeV2Test extends BaseStarTreeV2Test { @Override - ValueAggregator getValueAggregator() { + ValueAggregator getValueAggregator() { return new MinValueAggregator(); } @@ -39,7 +39,7 @@ DataType getRawValueType() { } @Override - Number getRandomRawValue(Random random) { + Object getRandomRawValue(Random random) { return random.nextFloat(); } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/SumPrecisionStarTreeV2Test.java b/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/SumPrecisionStarTreeV2Test.java index 87c74aa59152..df622a5e03b9 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/SumPrecisionStarTreeV2Test.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/SumPrecisionStarTreeV2Test.java @@ -41,7 +41,7 @@ DataType getRawValueType() { } @Override - Double getRandomRawValue(Random random) { + Object getRandomRawValue(Random random) { return random.nextDouble(); } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/SumStarTreeV2Test.java b/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/SumStarTreeV2Test.java index dc198c89ed65..e18449a20364 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/SumStarTreeV2Test.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/startree/v2/SumStarTreeV2Test.java @@ -26,10 +26,10 @@ import static org.testng.Assert.assertEquals; -public class SumStarTreeV2Test extends BaseStarTreeV2Test { +public class SumStarTreeV2Test extends BaseStarTreeV2Test { @Override - ValueAggregator getValueAggregator() { + ValueAggregator getValueAggregator() { return new SumValueAggregator(); } @@ -39,7 +39,7 @@ DataType getRawValueType() { } @Override - Number getRandomRawValue(Random random) { + Object getRandomRawValue(Random random) { return random.nextInt(); } diff --git a/pinot-core/src/test/resources/TableIndexingTest.csv b/pinot-core/src/test/resources/TableIndexingTest.csv index 8c39fe3bfaf5..420f2c72fe7a 100644 --- a/pinot-core/src/test/resources/TableIndexingTest.csv +++ b/pinot-core/src/test/resources/TableIndexingTest.csv @@ -344,7 +344,7 @@ STRING;sv;dict;json_index;false;Column: col Unrecognized token 'str': was expect STRING;sv;dict;native_text_index;true; STRING;sv;dict;text_index;true; STRING;sv;dict;range_index;true; -STRING;sv;dict;startree_index;false;class java.lang.String cannot be cast to class java.lang.Number (java.lang.String and java.lang.Number are in module java.base of loader 'bootstrap') +STRING;sv;dict;startree_index;false;For input string: "str-.*" STRING;sv;dict;vector_index;false;Cannot create vector index on single-value column: col STRING;sv;dict;multi_col_text_index;true; STRING;mv;dict;timestamp_index;false;Cannot create TIMESTAMP index on column: col of stored type other than LONG @@ -380,7 +380,7 @@ JSON;sv;dict;json_index;true; JSON;sv;dict;native_text_index;false;expected [1] but found [0] JSON;sv;dict;text_index;false;expected [1] but found [0] JSON;sv;dict;range_index;true; -JSON;sv;dict;startree_index;false;class java.lang.String cannot be cast to class java.lang.Number (java.lang.String and java.lang.Number are in module java.base of loader 'bootstrap') +JSON;sv;dict;startree_index;false;For input string: "\{"field":".*"\}" JSON;sv;dict;vector_index;false;Cannot create vector index on single-value column: col JSON;sv;dict;multi_col_text_index;false;Cannot create TEXT index on column: col of stored type other than STRING BYTES;sv;raw;timestamp_index;false;Cannot create TIMESTAMP index on column: col of stored type other than LONG diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/AvgValueAggregator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/AvgValueAggregator.java index 6a6124108c24..7ba199642e4a 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/AvgValueAggregator.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/AvgValueAggregator.java @@ -42,7 +42,7 @@ public AvgPair getInitialAggregatedValue(Object rawValue) { if (rawValue instanceof byte[]) { return deserializeAggregatedValue((byte[]) rawValue); } else { - return new AvgPair(((Number) rawValue).doubleValue(), 1L); + return new AvgPair(ValueAggregatorUtils.toDouble(rawValue), 1L); } } @@ -51,7 +51,7 @@ public AvgPair applyRawValue(AvgPair value, Object rawValue) { if (rawValue instanceof byte[]) { value.apply(deserializeAggregatedValue((byte[]) rawValue)); } else { - value.apply(((Number) rawValue).doubleValue(), 1L); + value.apply(ValueAggregatorUtils.toDouble(rawValue), 1L); } return value; } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/MaxValueAggregator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/MaxValueAggregator.java index 13669188731a..f5b51833faef 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/MaxValueAggregator.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/MaxValueAggregator.java @@ -22,7 +22,7 @@ import org.apache.pinot.spi.data.FieldSpec.DataType; -public class MaxValueAggregator implements ValueAggregator { +public class MaxValueAggregator implements ValueAggregator { public static final DataType AGGREGATED_VALUE_TYPE = DataType.DOUBLE; @Override @@ -36,13 +36,13 @@ public DataType getAggregatedValueType() { } @Override - public Double getInitialAggregatedValue(Number rawValue) { - return rawValue.doubleValue(); + public Double getInitialAggregatedValue(Object rawValue) { + return ValueAggregatorUtils.toDouble(rawValue); } @Override - public Double applyRawValue(Double value, Number rawValue) { - return Math.max(value, rawValue.doubleValue()); + public Double applyRawValue(Double value, Object rawValue) { + return Math.max(value, ValueAggregatorUtils.toDouble(rawValue)); } @Override diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/MinMaxRangeValueAggregator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/MinMaxRangeValueAggregator.java index 484e22b4f937..d6aeb1195ea9 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/MinMaxRangeValueAggregator.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/MinMaxRangeValueAggregator.java @@ -42,7 +42,7 @@ public MinMaxRangePair getInitialAggregatedValue(Object rawValue) { if (rawValue instanceof byte[]) { return deserializeAggregatedValue((byte[]) rawValue); } else { - double doubleValue = ((Number) rawValue).doubleValue(); + double doubleValue = ValueAggregatorUtils.toDouble(rawValue); return new MinMaxRangePair(doubleValue, doubleValue); } } @@ -52,8 +52,7 @@ public MinMaxRangePair applyRawValue(MinMaxRangePair value, Object rawValue) { if (rawValue instanceof byte[]) { value.apply(deserializeAggregatedValue((byte[]) rawValue)); } else { - double doubleValue = ((Number) rawValue).doubleValue(); - value.apply(doubleValue); + value.apply(ValueAggregatorUtils.toDouble(rawValue)); } return value; } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/MinValueAggregator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/MinValueAggregator.java index c7abc28979ea..2ebc95c123c7 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/MinValueAggregator.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/MinValueAggregator.java @@ -22,7 +22,7 @@ import org.apache.pinot.spi.data.FieldSpec.DataType; -public class MinValueAggregator implements ValueAggregator { +public class MinValueAggregator implements ValueAggregator { public static final DataType AGGREGATED_VALUE_TYPE = DataType.DOUBLE; @Override @@ -36,13 +36,13 @@ public DataType getAggregatedValueType() { } @Override - public Double getInitialAggregatedValue(Number rawValue) { - return rawValue.doubleValue(); + public Double getInitialAggregatedValue(Object rawValue) { + return ValueAggregatorUtils.toDouble(rawValue); } @Override - public Double applyRawValue(Double value, Number rawValue) { - return Math.min(value, rawValue.doubleValue()); + public Double applyRawValue(Double value, Object rawValue) { + return Math.min(value, ValueAggregatorUtils.toDouble(rawValue)); } @Override diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/PercentileEstValueAggregator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/PercentileEstValueAggregator.java index 7ff8a2415812..4915461b87e4 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/PercentileEstValueAggregator.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/PercentileEstValueAggregator.java @@ -51,7 +51,7 @@ public QuantileDigest getInitialAggregatedValue(Object rawValue) { _maxByteSize = Math.max(_maxByteSize, bytes.length); } else { initialValue = new QuantileDigest(DEFAULT_MAX_ERROR); - initialValue.add(((Number) rawValue).longValue()); + initialValue.add(toLong(rawValue)); _maxByteSize = Math.max(_maxByteSize, initialValue.getByteSize()); } return initialValue; @@ -62,12 +62,24 @@ public QuantileDigest applyRawValue(QuantileDigest value, Object rawValue) { if (rawValue instanceof byte[]) { value.merge(deserializeAggregatedValue((byte[]) rawValue)); } else { - value.add(((Number) rawValue).longValue()); + value.add(toLong(rawValue)); } _maxByteSize = Math.max(_maxByteSize, value.getByteSize()); return value; } + private static long toLong(Object rawValue) { + if (rawValue instanceof Number) { + return ((Number) rawValue).longValue(); + } + String stringValue = rawValue.toString(); + try { + return Long.parseLong(stringValue); + } catch (NumberFormatException e) { + return (long) Double.parseDouble(stringValue); + } + } + @Override public QuantileDigest applyAggregatedValue(QuantileDigest value, QuantileDigest aggregatedValue) { value.merge(aggregatedValue); diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/PercentileTDigestValueAggregator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/PercentileTDigestValueAggregator.java index da21438a4ed7..054b52f9bd76 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/PercentileTDigestValueAggregator.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/PercentileTDigestValueAggregator.java @@ -61,7 +61,7 @@ public TDigest getInitialAggregatedValue(Object rawValue) { _maxByteSize = Math.max(_maxByteSize, bytes.length); } else { initialValue = TDigest.createMergingDigest(_compressionFactor); - initialValue.add(((Number) rawValue).doubleValue()); + initialValue.add(ValueAggregatorUtils.toDouble(rawValue)); _maxByteSize = Math.max(_maxByteSize, initialValue.byteSize()); } return initialValue; @@ -72,7 +72,7 @@ public TDigest applyRawValue(TDigest value, Object rawValue) { if (rawValue instanceof byte[]) { value.add(deserializeAggregatedValue((byte[]) rawValue)); } else { - value.add(((Number) rawValue).doubleValue()); + value.add(ValueAggregatorUtils.toDouble(rawValue)); } _maxByteSize = Math.max(_maxByteSize, value.byteSize()); return value; diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/SumValueAggregator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/SumValueAggregator.java index 429b464c5489..1f79b6cc6c84 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/SumValueAggregator.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/SumValueAggregator.java @@ -22,7 +22,7 @@ import org.apache.pinot.spi.data.FieldSpec.DataType; -public class SumValueAggregator implements ValueAggregator { +public class SumValueAggregator implements ValueAggregator { public static final DataType AGGREGATED_VALUE_TYPE = DataType.DOUBLE; @Override @@ -36,13 +36,13 @@ public DataType getAggregatedValueType() { } @Override - public Double getInitialAggregatedValue(Number rawValue) { - return rawValue.doubleValue(); + public Double getInitialAggregatedValue(Object rawValue) { + return ValueAggregatorUtils.toDouble(rawValue); } @Override - public Double applyRawValue(Double value, Number rawValue) { - return value + rawValue.doubleValue(); + public Double applyRawValue(Double value, Object rawValue) { + return value + ValueAggregatorUtils.toDouble(rawValue); } @Override diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/ValueAggregatorUtils.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/ValueAggregatorUtils.java new file mode 100644 index 000000000000..3728e29baf6f --- /dev/null +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/ValueAggregatorUtils.java @@ -0,0 +1,36 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.segment.local.aggregator; + +public class ValueAggregatorUtils { + private ValueAggregatorUtils() { + } + + /// Tries to convert the given value to a double. + /// We need this for [ValueAggregator] because the raw value might not be converted to the desired data type yet if it + /// is not specified in the schema. + /// TODO: Provide a way to specify the desired data type for the raw value. + public static double toDouble(Object value) { + if (value instanceof Number) { + return ((Number) value).doubleValue(); + } else { + return Double.parseDouble(value.toString()); + } + } +} From e50947a53f9eec910d3f7a8aa80b082c5f6367fd Mon Sep 17 00:00:00 2001 From: Rajat Venkatesh <1638298+vrajat@users.noreply.github.com> Date: Wed, 6 Aug 2025 16:32:18 +0530 Subject: [PATCH 086/167] Add tests for OOM Termination of MSE queries. (#16514) --- ...ClusterMemBasedServerQueryKillingTest.java | 138 ++++++++++++++---- 1 file changed, 113 insertions(+), 25 deletions(-) diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterMemBasedServerQueryKillingTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterMemBasedServerQueryKillingTest.java index b39cb2375049..efa34e69eb57 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterMemBasedServerQueryKillingTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterMemBasedServerQueryKillingTest.java @@ -37,7 +37,9 @@ import org.apache.log4j.Level; import org.apache.log4j.LogManager; import org.apache.pinot.core.accounting.PerQueryCPUMemAccountantFactory; +import org.apache.pinot.spi.accounting.QueryResourceTracker; import org.apache.pinot.spi.accounting.ThreadResourceUsageProvider; +import org.apache.pinot.spi.config.instance.InstanceType; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.data.FieldSpec; @@ -48,6 +50,7 @@ import org.apache.pinot.spi.utils.CommonConstants; import org.apache.pinot.spi.utils.builder.TableConfigBuilder; import org.apache.pinot.util.TestUtils; +import org.testng.annotations.AfterMethod; import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; @@ -100,6 +103,57 @@ protected int getNumServers() { return NUM_SERVERS; } + /** + * Keeps track of metadata of queries that have been terminated due to OOM. + * This is used to verify that the query was killed and the method used to kill the query. + */ + public static class TestAccountantFactory extends PerQueryCPUMemAccountantFactory { + @Override + public PerQueryCPUMemResourceUsageAccountant init(PinotConfiguration config, String instanceId, + InstanceType instanceType) { + return new TestAccountant(config, instanceId, instanceType); + } + + public static class TestAccountant extends PerQueryCPUMemAccountantFactory.PerQueryCPUMemResourceUsageAccountant { + private QueryResourceTracker _queryResourceTracker = null; + private long _totalHeapMemoryUsage = 0L; + private boolean _hasCallback = false; + + public TestAccountant(PinotConfiguration config, String instanceId, InstanceType instanceType) { + super(config, instanceId, instanceType); + } + + @Override + public void logTerminatedQuery(QueryResourceTracker queryResourceTracker, long totalHeapMemoryUsage, + boolean hasCallback) { + super.logTerminatedQuery(queryResourceTracker, totalHeapMemoryUsage, hasCallback); + _queryResourceTracker = queryResourceTracker; + _totalHeapMemoryUsage = totalHeapMemoryUsage; + _hasCallback = hasCallback; + } + + public void reset() { + _queryResourceTracker = null; + _totalHeapMemoryUsage = 0L; + _hasCallback = false; + } + + public QueryResourceTracker getQueryResourceTracker() { + return _queryResourceTracker; + } + + public long getTotalHeapMemoryUsage() { + return _totalHeapMemoryUsage; + } + + public boolean hasCallback() { + return _hasCallback; + } + } + } + + private TestAccountantFactory.TestAccountant _testAccountant = null; + @BeforeClass public void setUp() throws Exception { @@ -114,9 +168,9 @@ public void setUp() startZk(); startController(); startServers(); - while (!Tracing.isAccountantRegistered()) { - Thread.sleep(100L); - } + TestUtils.waitForCondition(aVoid -> Tracing.isAccountantRegistered(), 100L, 60_000L, + "Waiting for accountant to be registered"); + _testAccountant = (TestAccountantFactory.TestAccountant) Tracing.getThreadAccountant(); startBrokers(); @@ -158,6 +212,13 @@ protected void startServers() startServers(getNumServers()); } + @AfterMethod + public void resetTestAccountant() { + if (_testAccountant != null) { + _testAccountant.reset(); + } + } + protected void overrideServerConf(PinotConfiguration serverConf) { serverConf.setProperty(CommonConstants.PINOT_QUERY_SCHEDULER_PREFIX + "." + CommonConstants.Accounting.CONFIG_OF_ALARMING_LEVEL_HEAP_USAGE_RATIO, 0.0f); @@ -165,7 +226,7 @@ protected void overrideServerConf(PinotConfiguration serverConf) { + CommonConstants.Accounting.CONFIG_OF_CRITICAL_LEVEL_HEAP_USAGE_RATIO, 0.15f); serverConf.setProperty( CommonConstants.PINOT_QUERY_SCHEDULER_PREFIX + "." + CommonConstants.Accounting.CONFIG_OF_FACTORY_NAME, - "org.apache.pinot.core.accounting.PerQueryCPUMemAccountantFactory"); + TestAccountantFactory.class.getName()); serverConf.setProperty(CommonConstants.PINOT_QUERY_SCHEDULER_PREFIX + "." + CommonConstants.Accounting.CONFIG_OF_ENABLE_THREAD_MEMORY_SAMPLING, true); serverConf.setProperty( @@ -185,7 +246,7 @@ protected void overrideBrokerConf(PinotConfiguration brokerConf) { + CommonConstants.Accounting.CONFIG_OF_PANIC_LEVEL_HEAP_USAGE_RATIO, 1.1f); brokerConf.setProperty(CommonConstants.PINOT_QUERY_SCHEDULER_PREFIX + "." + CommonConstants.Accounting.CONFIG_OF_FACTORY_NAME, - "org.apache.pinot.core.accounting.PerQueryCPUMemAccountantFactory"); + TestAccountantFactory.class.getName()); brokerConf.setProperty(CommonConstants.PINOT_QUERY_SCHEDULER_PREFIX + "." + CommonConstants.Accounting.CONFIG_OF_ENABLE_THREAD_MEMORY_SAMPLING, true); brokerConf.setProperty(CommonConstants.PINOT_QUERY_SCHEDULER_PREFIX + "." @@ -214,22 +275,30 @@ protected TableConfig createOfflineTableConfig() { .build(); } - @Test(dataProvider = "useBothQueryEngines") - public void testDigestOOM(boolean useMultiStageQueryEngine) + @Test + public void testDigestOOM() throws Exception { - setUseMultiStageQueryEngine(useMultiStageQueryEngine); - notSupportedInV2(); JsonNode queryResponse = postQuery(OOM_QUERY); String exceptionsNode = queryResponse.get("exceptions").toString(); assertTrue(exceptionsNode.contains("\"errorCode\":" + QueryErrorCode.QUERY_CANCELLATION.getId()), exceptionsNode); assertTrue(exceptionsNode.contains("got killed because"), exceptionsNode); } - @Test(dataProvider = "useBothQueryEngines") - public void testMemoryAllocationStats(boolean useMultiStageQueryEngine) + @Test + public void testDigestOOMMSE() + throws Exception { + setUseMultiStageQueryEngine(true); + JsonNode queryResponse = postQuery(OOM_QUERY); + String exceptionsNode = queryResponse.get("exceptions").toString(); + assertTrue(exceptionsNode.contains("\"errorCode\":" + QueryErrorCode.INTERNAL.getId()), exceptionsNode); + assertTrue(exceptionsNode.contains("Received 1 error from stage 1"), exceptionsNode); + assertTrue(_testAccountant.hasCallback()); + assertEquals(queryResponse.get("requestId").asText(), _testAccountant.getQueryResourceTracker().getQueryId()); + } + + @Test + public void testMemoryAllocationStats() throws Exception { - setUseMultiStageQueryEngine(useMultiStageQueryEngine); - notSupportedInV2(); JsonNode queryResponse = postQuery(COUNT_STAR_QUERY); long offlineThreadMemAllocatedBytes = queryResponse.get("offlineThreadMemAllocatedBytes").asLong(); long offlineResponseSerMemAllocatedBytes = queryResponse.get("offlineResponseSerMemAllocatedBytes").asLong(); @@ -240,11 +309,9 @@ public void testMemoryAllocationStats(boolean useMultiStageQueryEngine) assertEquals(offlineThreadMemAllocatedBytes + offlineResponseSerMemAllocatedBytes, offlineTotalMemAllocatedBytes); } - @Test(dataProvider = "useBothQueryEngines") - public void testSelectionOnlyOOM(boolean useMultiStageQueryEngine) + @Test + public void testSelectionOnlyOOM() throws Exception { - setUseMultiStageQueryEngine(useMultiStageQueryEngine); - notSupportedInV2(); JsonNode queryResponse = postQuery(OOM_QUERY_SELECTION_ONLY); String exceptionsNode = queryResponse.get("exceptions").toString(); @@ -252,21 +319,42 @@ public void testSelectionOnlyOOM(boolean useMultiStageQueryEngine) assertTrue(exceptionsNode.contains("got killed because"), exceptionsNode); } - @Test(dataProvider = "useBothQueryEngines") - public void testDigestOOM2(boolean useMultiStageQueryEngine) + @Test + public void testSelectionOnlyOOMMSE() + throws Exception { + setUseMultiStageQueryEngine(true); + JsonNode queryResponse = postQuery(OOM_QUERY_SELECTION_ONLY); + + String exceptionsNode = queryResponse.get("exceptions").toString(); + assertTrue(exceptionsNode.contains("\"errorCode\":" + QueryErrorCode.INTERNAL.getId()), exceptionsNode); + assertTrue(exceptionsNode.contains("Received 1 error from stage 1"), exceptionsNode); + assertTrue(_testAccountant.hasCallback()); + assertEquals(queryResponse.get("requestId").asText(), _testAccountant.getQueryResourceTracker().getQueryId()); + } + + @Test + public void testDigestOOM2() throws Exception { - setUseMultiStageQueryEngine(useMultiStageQueryEngine); - notSupportedInV2(); JsonNode queryResponse = postQuery(OOM_QUERY_2); String exceptionsNode = queryResponse.get("exceptions").toString(); assertTrue(exceptionsNode.contains("got killed because"), exceptionsNode); } - @Test(dataProvider = "useBothQueryEngines") - public void testDigestOOMMultipleQueries(boolean useMultiStageQueryEngine) + @Test + public void testDigestOOM2MSE() + throws Exception { + setUseMultiStageQueryEngine(true); + JsonNode queryResponse = postQuery(OOM_QUERY_2); + String exceptionsNode = queryResponse.get("exceptions").toString(); + assertTrue(exceptionsNode.contains("\"errorCode\":" + QueryErrorCode.INTERNAL.getId()), exceptionsNode); + assertTrue(exceptionsNode.contains("Received 1 error from stage 1"), exceptionsNode); + assertTrue(_testAccountant.hasCallback()); + assertEquals(queryResponse.get("requestId").asText(), _testAccountant.getQueryResourceTracker().getQueryId()); + } + + @Test + public void testDigestOOMMultipleQueries() throws Exception { - setUseMultiStageQueryEngine(useMultiStageQueryEngine); - notSupportedInV2(); AtomicReference queryResponse1 = new AtomicReference<>(); AtomicReference queryResponse2 = new AtomicReference<>(); AtomicReference queryResponse3 = new AtomicReference<>(); From ad4a39f6ede88788e558891c0ea2b0d76988ba50 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Wed, 6 Aug 2025 09:55:30 -0700 Subject: [PATCH 087/167] Bump com.nimbusds:nimbus-jose-jwt from 10.4 to 10.4.1 (#16522) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index a284af85d8be..ea535c25e099 100644 --- a/pom.xml +++ b/pom.xml @@ -257,7 +257,7 @@ 3.30.4 2.0.1 1.5.4 - 10.4 + 10.4.1 3.6.3 9.4.57.v20241219 7.1.1 From c4ea4ea7a3cb1e956ae6dee8b846994c79d9b4f9 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Wed, 6 Aug 2025 09:58:15 -0700 Subject: [PATCH 088/167] Bump software.amazon.awssdk:bom from 2.32.15 to 2.32.16 (#16523) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index ea535c25e099..9d12ed3a194c 100644 --- a/pom.xml +++ b/pom.xml @@ -178,7 +178,7 @@ 0.15.1 0.4.7 4.2.2 - 2.32.15 + 2.32.16 1.2.37 1.22.0 2.14.0 From 021368c5374b93a980c5b6fee36ab2b0809e6091 Mon Sep 17 00:00:00 2001 From: Yash Mayya Date: Wed, 6 Aug 2025 22:29:13 +0530 Subject: [PATCH 089/167] Fix typo in semi-join distinct project optimization rule (#16520) --- .../pinot/calcite/rel/rules/PinotQueryRuleSets.java | 4 ++-- ...le.java => PinotSemiJoinDistinctProjectRule.java} | 12 ++++++------ .../org/apache/pinot/spi/utils/CommonConstants.java | 2 +- 3 files changed, 9 insertions(+), 9 deletions(-) rename pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/{PinotSeminJoinDistinctProjectRule.java => PinotSemiJoinDistinctProjectRule.java} (85%) diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java index e07de5a53e9a..b9d2604b83e2 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotQueryRuleSets.java @@ -116,8 +116,8 @@ private PinotQueryRuleSets() { // join and semi-join rules SemiJoinRule.ProjectToSemiJoinRule.ProjectToSemiJoinRuleConfig.DEFAULT .withDescription(PlannerRuleNames.PROJECT_TO_SEMI_JOIN).toRule(), - PinotSeminJoinDistinctProjectRule - .instanceWithDescription(PlannerRuleNames.SEMIN_JOIN_DISTINCT_PROJECT), + PinotSemiJoinDistinctProjectRule + .instanceWithDescription(PlannerRuleNames.SEMI_JOIN_DISTINCT_PROJECT), // Consider semijoin optimizations first before push transitive predicate // Pinot version doesn't push predicates to the right in case of lookup join diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotSeminJoinDistinctProjectRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotSemiJoinDistinctProjectRule.java similarity index 85% rename from pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotSeminJoinDistinctProjectRule.java rename to pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotSemiJoinDistinctProjectRule.java index 95f25b7f78e9..0b1766fe704a 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotSeminJoinDistinctProjectRule.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/calcite/rel/rules/PinotSemiJoinDistinctProjectRule.java @@ -38,15 +38,15 @@ * {@link org.apache.calcite.rel.logical.LogicalProject} on top of a Semi join * {@link org.apache.calcite.rel.core.Join} to ensure the correctness of the query. */ -public class PinotSeminJoinDistinctProjectRule extends RelOptRule { - public static final PinotSeminJoinDistinctProjectRule INSTANCE = - new PinotSeminJoinDistinctProjectRule(PinotRuleUtils.PINOT_REL_FACTORY, null); +public class PinotSemiJoinDistinctProjectRule extends RelOptRule { + public static final PinotSemiJoinDistinctProjectRule INSTANCE = + new PinotSemiJoinDistinctProjectRule(PinotRuleUtils.PINOT_REL_FACTORY, null); - public static PinotSeminJoinDistinctProjectRule instanceWithDescription(String description) { - return new PinotSeminJoinDistinctProjectRule(PinotRuleUtils.PINOT_REL_FACTORY, description); + public static PinotSemiJoinDistinctProjectRule instanceWithDescription(String description) { + return new PinotSemiJoinDistinctProjectRule(PinotRuleUtils.PINOT_REL_FACTORY, description); } - public PinotSeminJoinDistinctProjectRule(RelBuilderFactory factory, @Nullable String description) { + public PinotSemiJoinDistinctProjectRule(RelBuilderFactory factory, @Nullable String description) { super(operand(LogicalJoin.class, operand(AbstractRelNode.class, any()), operand(LogicalProject.class, any())), factory, description); } diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java index 894c73f8dd91..367c93704465 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java @@ -783,7 +783,7 @@ public static class PlannerRuleNames { public static final String EVALUATE_LITERAL_FILTER = "EvaluateFilterLiteral"; public static final String JOIN_PUSH_EXPRESSIONS = "JoinPushExpressions"; public static final String PROJECT_TO_SEMI_JOIN = "ProjectToSemiJoin"; - public static final String SEMIN_JOIN_DISTINCT_PROJECT = "SeminJoinDistinctProject"; + public static final String SEMI_JOIN_DISTINCT_PROJECT = "SemiJoinDistinctProject"; public static final String UNION_TO_DISTINCT = "UnionToDistinct"; public static final String AGGREGATE_REMOVE = "AggregateRemove"; public static final String AGGREGATE_JOIN_TRANSPOSE = "AggregateJoinTranspose"; From d5206cb7c4035f8a623986bc19176ca2cb0a9fd7 Mon Sep 17 00:00:00 2001 From: "Xiaotian (Jackie) Jiang" <17555551+Jackie-Jiang@users.noreply.github.com> Date: Wed, 6 Aug 2025 14:17:34 -0600 Subject: [PATCH 090/167] Remove ResourceManagerAccountingTest.testWorkerThreadInterruption() (#16525) --- .../ResourceManagerAccountingTest.java | 54 ------------------- 1 file changed, 54 deletions(-) diff --git a/pinot-core/src/test/java/org/apache/pinot/core/accounting/ResourceManagerAccountingTest.java b/pinot-core/src/test/java/org/apache/pinot/core/accounting/ResourceManagerAccountingTest.java index 2641cd60bd3b..89b2dbc27372 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/accounting/ResourceManagerAccountingTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/accounting/ResourceManagerAccountingTest.java @@ -20,18 +20,15 @@ import java.io.File; import java.io.IOException; -import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Random; import java.util.concurrent.CountDownLatch; -import java.util.concurrent.ExecutionException; import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicInteger; -import java.util.concurrent.atomic.AtomicReference; import java.util.function.Supplier; import org.apache.commons.io.FileUtils; import org.apache.log4j.Level; @@ -298,57 +295,6 @@ public void testWorkloadLevelThreadMemoryAccounting() Thread.sleep(1000_000); } - - /** - * Test the mechanism of worker thread checking for runnerThread's interruption flag - */ - @Test - public void testWorkerThreadInterruption() - throws Exception { - ResourceManager rm = getResourceManager(2, 5, 1, 3, Collections.emptyMap(), - new Tracing.DefaultThreadResourceUsageAccountant()); - AtomicReference[] futures = new AtomicReference[5]; - for (int i = 0; i < 5; i++) { - futures[i] = new AtomicReference<>(); - } - ThreadResourceUsageProvider.setThreadCpuTimeMeasurementEnabled(true); - AtomicReference runnerThread = new AtomicReference<>(); - rm.getQueryRunners().submit(() -> { - Thread thread = Thread.currentThread(); - runnerThread.set(thread); - for (int j = 0; j < 5; j++) { - futures[j].set(rm.getQueryWorkers().submit(() -> { - for (int i = 0; i < 1000000; i++) { - try { - Thread.sleep(5); - } catch (InterruptedException ignored) { - } - if (thread.isInterrupted()) { - throw new EarlyTerminationException(); - } - } - })); - } - while (true) { - } - }); - Thread.sleep(50); - runnerThread.get().interrupt(); - - for (int i = 0; i < 5; i++) { - try { - futures[i].get().get(); - } catch (ExecutionException e) { - Assert.assertFalse(futures[i].get().isCancelled()); - Assert.assertTrue(futures[i].get().isDone()); - Assert.assertTrue(e.getMessage().contains("EarlyTerminationException"), - "Error message should contain EarlyTerminationException, found: " + e.getMessage()); - return; - } - } - Assert.fail("Expected EarlyTerminationException to be thrown"); - } - /** * Test instrumentation during {@link DataTable} creation */ From 4b83b8bb26d5fb13729ffa29a0ad7a55ed8a4e3a Mon Sep 17 00:00:00 2001 From: Gonzalo Ortiz Jaureguizar Date: Wed, 6 Aug 2025 22:26:54 +0200 Subject: [PATCH 091/167] Fix SSE subquery handling and clarify query start semantics (#16524) --- .../broker/requesthandler/BaseBrokerRequestHandler.java | 6 ++++++ .../requesthandler/BaseSingleStageBrokerRequestHandler.java | 2 +- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java index 97b87ce26f7c..59498218e254 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java @@ -394,6 +394,12 @@ protected String extractClientRequestId(SqlNodeAndOptions sqlNodeAndOptions) { ? sqlNodeAndOptions.getOptions().get(Broker.Request.QueryOptionKey.CLIENT_QUERY_ID) : null; } + /** + * Called when a query starts + * TODO: This method was created to keep track of running queries for cancellation, but it is useful for other uses. + * But right now the semantics are not clear. For example, while MSE calls this method once, SSE calls it once per + * query AND subquery, which means this method is called multiple times for the same query. + */ protected void onQueryStart(long requestId, String clientRequestId, String query, Object... extras) { if (isQueryCancellationEnabled()) { _queriesById.put(requestId, query); diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java index 6aac44b91cb9..d405e998b681 100644 --- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java +++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseSingleStageBrokerRequestHandler.java @@ -321,6 +321,7 @@ protected BrokerResponse handleRequest(long requestId, String query, SqlNodeAndO JsonNode request, @Nullable RequesterIdentity requesterIdentity, RequestContext requestContext, @Nullable HttpHeaders httpHeaders, AccessControl accessControl) throws Exception { + QueryThreadContext.setQueryEngine("sse"); _queryLogger.log(requestId, query); //Start instrumentation context. This must not be moved further below interspersed into the code. @@ -1042,7 +1043,6 @@ static String addRoutingPolicyInErrMsg(String errorMessage, String realtimeRouti @Override protected void onQueryStart(long requestId, String clientRequestId, String query, Object... extras) { super.onQueryStart(requestId, clientRequestId, query, extras); - QueryThreadContext.setQueryEngine("sse"); if (isQueryCancellationEnabled() && extras.length > 0 && extras[0] instanceof QueryServers) { _serversById.put(requestId, (QueryServers) extras[0]); } From 78dabced284493933d526e2fce61f239443bf5f1 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Wed, 6 Aug 2025 13:58:14 -0700 Subject: [PATCH 092/167] feat: Enhance Controller UI with comprehensive minion task stats and status filtering (#16521) * add more minion sub task stats * feat: Add status filtering for minion tasks and subtasks - Created TaskStatusFilter component with visual status chips - Added status filtering to task listing pages with dropdown selection - Enhanced TaskDetail component with subtask status filtering - Implemented comprehensive filtering for all task states: COMPLETED, RUNNING, WAITING, ERROR, UNKNOWN, DROPPED, TIMED_OUT, ABORTED - Optimized React hooks with useCallback to prevent unnecessary re-renders - Added proper TypeScript types and ESLint compliance - Improved user experience with color-coded status indicators --- .../app/components/TaskStatusFilter.tsx | 215 ++++++++++++++++++ .../app/components/useTaskListing.tsx | 49 +++- .../main/resources/app/pages/TaskDetail.tsx | 50 +++- .../resources/app/utils/PinotMethodUtils.ts | 16 +- 4 files changed, 317 insertions(+), 13 deletions(-) create mode 100644 pinot-controller/src/main/resources/app/components/TaskStatusFilter.tsx diff --git a/pinot-controller/src/main/resources/app/components/TaskStatusFilter.tsx b/pinot-controller/src/main/resources/app/components/TaskStatusFilter.tsx new file mode 100644 index 000000000000..039952a5fa53 --- /dev/null +++ b/pinot-controller/src/main/resources/app/components/TaskStatusFilter.tsx @@ -0,0 +1,215 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +import React from 'react'; +import { + FormControl, + InputLabel, + Select, + MenuItem, + makeStyles, + Chip +} from '@material-ui/core'; + +const useStyles = makeStyles((theme) => ({ + formControl: { + minWidth: 180, + marginRight: theme.spacing(2), + }, + inputLabel: { + backgroundColor: 'white', + paddingLeft: theme.spacing(1), + paddingRight: theme.spacing(1), + color: theme.palette.text.secondary, + }, + select: { + '& .MuiSelect-select': { + padding: '8px 12px', + }, + }, + menuPaper: { + maxHeight: 300, + marginTop: 4, + }, + menuItem: { + padding: '8px 16px', + '&:hover': { + backgroundColor: theme.palette.action.hover, + }, + }, + statusChip: { + fontSize: '0.75rem', + height: 20, + borderRadius: 10, + fontWeight: 500, + }, + completed: { + backgroundColor: '#e8f5e8', + borderColor: '#4caf50', + color: '#2e7d32', + }, + running: { + backgroundColor: '#e3f2fd', + borderColor: '#2196f3', + color: '#1565c0', + }, + waiting: { + backgroundColor: '#fff3e0', + borderColor: '#ff9800', + color: '#ef6c00', + }, + error: { + backgroundColor: '#ffebee', + borderColor: '#f44336', + color: '#c62828', + }, + unknown: { + backgroundColor: '#f3e5f5', + borderColor: '#9c27b0', + color: '#7b1fa2', + }, + dropped: { + backgroundColor: '#fce4ec', + borderColor: '#e91e63', + color: '#ad1457', + }, + timedout: { + backgroundColor: '#fff8e1', + borderColor: '#ffc107', + color: '#f57c00', + }, + aborted: { + backgroundColor: '#efebe9', + borderColor: '#795548', + color: '#5d4037', + }, +})); + +export type TaskStatus = 'COMPLETED' | 'RUNNING' | 'WAITING' | 'ERROR' | 'UNKNOWN' | 'DROPPED' | 'TIMED_OUT' | 'ABORTED'; + +type TaskStatusFilterOption = { + label: string; + value: 'ALL' | TaskStatus; +}; + +type TaskStatusFilterProps = { + value: 'ALL' | TaskStatus; + onChange: (value: 'ALL' | TaskStatus) => void; + options: TaskStatusFilterOption[]; +}; + +export const getTaskStatusChipClass = (status: string, classes?: any) => { + if (!classes) return ''; + + switch (status.toUpperCase()) { + case 'COMPLETED': + return classes.completed; + case 'RUNNING': + return classes.running; + case 'WAITING': + return classes.waiting; + case 'ERROR': + return classes.error; + case 'UNKNOWN': + return classes.unknown; + case 'DROPPED': + return classes.dropped; + case 'TIMED_OUT': + case 'TIMEDOUT': + return classes.timedout; + case 'ABORTED': + return classes.aborted; + default: + return ''; + } +}; + +const TaskStatusFilter: React.FC = ({ value, onChange, options }) => { + const classes = useStyles(); + + const renderValue = (selected: string) => { + const selectedOption = options.find(option => option.value === selected); + const label = selectedOption ? selectedOption.label : 'All'; + + if (selected === 'ALL') { + return label; + } + + return ( +

    + +
    + ); + }; + + return ( + + Status Filter + + + ); +}; + +export default TaskStatusFilter; diff --git a/pinot-controller/src/main/resources/app/components/useTaskListing.tsx b/pinot-controller/src/main/resources/app/components/useTaskListing.tsx index 5e6d6f6972d1..4eaad35e2a9e 100644 --- a/pinot-controller/src/main/resources/app/components/useTaskListing.tsx +++ b/pinot-controller/src/main/resources/app/components/useTaskListing.tsx @@ -17,9 +17,11 @@ * under the License. */ -import React, { useEffect, useState } from 'react'; +import React, { useEffect, useState, useMemo, useCallback } from 'react'; +import { Box } from '@material-ui/core'; import { TableData } from 'Models'; import CustomizedTables from './Table'; +import TaskStatusFilter, { TaskStatus } from './TaskStatusFilter'; import PinotMethodUtils from '../utils/PinotMethodUtils'; import { useTimezone } from '../contexts/TimezoneContext'; @@ -28,17 +30,53 @@ export default function useTaskListing(props) { const { currentTimezone } = useTimezone(); const [fetching, setFetching] = useState(true); const [tasks, setTasks] = useState({ records: [], columns: [] }); + const [statusFilter, setStatusFilter] = useState<'ALL' | TaskStatus>('ALL'); - const fetchData = async () => { + const fetchData = useCallback(async () => { setFetching(true); const tasksRes = await PinotMethodUtils.getTasksList(tableName, taskType); setTasks(tasksRes); setFetching(false); - }; + }, [tableName, taskType]); useEffect(() => { fetchData(); - }, [currentTimezone]); + }, [currentTimezone, fetchData]); + + const filteredTasks = useMemo(() => { + if (statusFilter === 'ALL') { + return tasks; + } + + const filtered = tasks.records.filter(([_, status]) => { + const taskStatus = typeof status === 'object' && status !== null && 'value' in status + ? status.value as string + : status as string; + return taskStatus.toUpperCase() === statusFilter; + }); + + return { ...tasks, records: filtered }; + }, [tasks, statusFilter]); + + const statusFilterOptions = [ + { label: 'All', value: 'ALL' as const }, + { label: 'Completed', value: 'COMPLETED' as const }, + { label: 'Running', value: 'RUNNING' as const }, + { label: 'Waiting', value: 'WAITING' as const }, + { label: 'Error', value: 'ERROR' as const }, + { label: 'Unknown', value: 'UNKNOWN' as const }, + { label: 'Dropped', value: 'DROPPED' as const }, + { label: 'Timed Out', value: 'TIMED_OUT' as const }, + { label: 'Aborted', value: 'ABORTED' as const }, + ]; + + const statusFilterElement = ( + + ); return { tasks, @@ -46,11 +84,12 @@ export default function useTaskListing(props) { content: !fetching && ( {statusFilterElement}} /> ) }; diff --git a/pinot-controller/src/main/resources/app/pages/TaskDetail.tsx b/pinot-controller/src/main/resources/app/pages/TaskDetail.tsx index 94468a3982ee..3d779e49b3f1 100644 --- a/pinot-controller/src/main/resources/app/pages/TaskDetail.tsx +++ b/pinot-controller/src/main/resources/app/pages/TaskDetail.tsx @@ -17,11 +17,12 @@ * under the License. */ -import React, { useEffect, useState } from 'react'; +import React, { useEffect, useState, useMemo, useCallback } from 'react'; import { get, each } from 'lodash'; -import { Grid, makeStyles } from '@material-ui/core'; +import { Grid, makeStyles, Box } from '@material-ui/core'; import PinotMethodUtils from '../utils/PinotMethodUtils'; import CustomizedTables from '../components/Table'; +import TaskStatusFilter, { TaskStatus } from '../components/TaskStatusFilter'; import { TaskRuntimeConfig } from 'Models'; import AppLoader from '../components/AppLoader'; import SimpleAccordion from '../components/SimpleAccordion'; @@ -77,8 +78,9 @@ const TaskDetail = (props) => { const [taskDebugData, setTaskDebugData] = useState({}); const [subtaskTableData, setSubtaskTableData] = useState({ columns: ['Task ID', 'Status', 'Start Time', 'Finish Time', 'Minion Host Name'], records: [] }); const [taskRuntimeConfig, setTaskRuntimeConfig] = useState(null); + const [subtaskStatusFilter, setSubtaskStatusFilter] = useState<'ALL' | TaskStatus>('ALL'); - const fetchData = async () => { + const fetchData = useCallback(async () => { setFetching(true); const [debugRes, runtimeConfig] = await Promise.all([ PinotMethodUtils.getTaskDebugData(taskID), @@ -101,11 +103,46 @@ const TaskDetail = (props) => { setTaskRuntimeConfig(runtimeConfig) setFetching(false); - }; + }, [taskID]); useEffect(() => { fetchData(); - }, [currentTimezone]); + }, [currentTimezone, fetchData]); + + const filteredSubtaskTableData = useMemo(() => { + if (subtaskStatusFilter === 'ALL') { + return subtaskTableData; + } + + const filtered = subtaskTableData.records.filter(([_, status]) => { + const subtaskStatus = typeof status === 'object' && status !== null && 'value' in status + ? status.value as string + : status as string; + return subtaskStatus.toUpperCase() === subtaskStatusFilter; + }); + + return { ...subtaskTableData, records: filtered }; + }, [subtaskTableData, subtaskStatusFilter]); + + const subtaskStatusFilterOptions = [ + { label: 'All', value: 'ALL' as const }, + { label: 'Completed', value: 'COMPLETED' as const }, + { label: 'Running', value: 'RUNNING' as const }, + { label: 'Waiting', value: 'WAITING' as const }, + { label: 'Error', value: 'ERROR' as const }, + { label: 'Unknown', value: 'UNKNOWN' as const }, + { label: 'Dropped', value: 'DROPPED' as const }, + { label: 'Timed Out', value: 'TIMED_OUT' as const }, + { label: 'Aborted', value: 'ABORTED' as const }, + ]; + + const subtaskStatusFilterElement = ( + + ); if(fetching) { return @@ -157,11 +194,12 @@ const TaskDetail = (props) => { {subtaskStatusFilterElement}} />     diff --git a/pinot-controller/src/main/resources/app/utils/PinotMethodUtils.ts b/pinot-controller/src/main/resources/app/utils/PinotMethodUtils.ts index e9428b3ad7d0..6fcd6a819267 100644 --- a/pinot-controller/src/main/resources/app/utils/PinotMethodUtils.ts +++ b/pinot-controller/src/main/resources/app/utils/PinotMethodUtils.ts @@ -934,7 +934,7 @@ const getElapsedTime = (startTime) => { const getTasksList = async (tableName, taskType) => { const { formatTimeInTimezone } = await import('./TimezoneUtils'); const finalResponse = { - columns: ['Task ID', 'Status', 'Start Time', 'Finish Time', 'Num of Sub Tasks'], + columns: ['Task ID', 'Status', 'Start Time', 'Finish Time', 'Sub Tasks (Total/Completed/Running/Waiting/Error/Other)'], records: [] } await new Promise((resolve, reject) => { @@ -942,12 +942,24 @@ const getTasksList = async (tableName, taskType) => { const promiseArr = []; const fetchInfo = async (taskID, status) => { const debugData = await getTaskDebugData(taskID); + const subtaskCount = get(debugData, 'data.subtaskCount', {}); + const total = get(subtaskCount, 'total', 0); + const completed = get(subtaskCount, 'completed', 0); + const running = get(subtaskCount, 'running', 0); + const waiting = get(subtaskCount, 'waiting', 0); + const error = get(subtaskCount, 'error', 0); + const unknown = get(subtaskCount, 'unknown', 0); + const dropped = get(subtaskCount, 'dropped', 0); + const timedOut = get(subtaskCount, 'timedOut', 0); + const aborted = get(subtaskCount, 'aborted', 0); + const other = unknown + dropped + timedOut + aborted; + finalResponse.records.push([ taskID, status, get(debugData, 'data.startTime') ? formatTimeInTimezone(get(debugData, 'data.startTime'), 'MMMM Do YYYY, HH:mm:ss z') : '', get(debugData, 'data.finishTime') ? formatTimeInTimezone(get(debugData, 'data.finishTime'), 'MMMM Do YYYY, HH:mm:ss z') : '', - get(debugData, 'data.subtaskCount.total', 0) + `${total}/${completed}/${running}/${waiting}/${error}/${other}` ]); }; each(response.data, async (val, key) => { From 4323a72887260205f6e00029d0989abb2a912cad Mon Sep 17 00:00:00 2001 From: "Xiaotian (Jackie) Jiang" <17555551+Jackie-Jiang@users.noreply.github.com> Date: Wed, 6 Aug 2025 17:43:11 -0600 Subject: [PATCH 093/167] Enhance StringFunctions.replace() (#16528) --- .../scalar/string/StringFunctions.java | 30 ++++++++----------- 1 file changed, 12 insertions(+), 18 deletions(-) diff --git a/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/string/StringFunctions.java b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/string/StringFunctions.java index 75bad424f113..4255ea7a2c6e 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/string/StringFunctions.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/string/StringFunctions.java @@ -19,6 +19,7 @@ package org.apache.pinot.common.function.scalar.string; import org.apache.commons.lang3.StringUtils; +import org.apache.commons.lang3.Strings; import org.apache.pinot.spi.annotations.ScalarFunction; @@ -54,35 +55,28 @@ public String concatWS(String separator, String input1, String input2) { } /** - * @param input - * @param searchString target substring to replace - * @param substitute new substring to be replaced with target - * @see String#replaceAll(String, String) + * @see Strings#replace(String, String, String) */ @ScalarFunction - public String replace(String input, String searchString, String substitute) { - if (StringUtils.isEmpty(input) || StringUtils.isEmpty(searchString) || substitute == null) { - return input; + public String replace(String text, String searchString, String replacement) { + if (text.isEmpty() || searchString.isEmpty()) { + return text; } int start = 0; - int end = StringUtils.indexOf(input, searchString, start); + int end = Strings.CS.indexOf(text, searchString, start); if (end == StringUtils.INDEX_NOT_FOUND) { - return input; + return text; } final int replLength = searchString.length(); - int increase = Math.max(substitute.length() - replLength, 0) * 16; + int increase = Math.max(replacement.length() - replLength, 0) * 16; _buffer.setLength(0); - _buffer.ensureCapacity(input.length() + increase); - int max = -1; + _buffer.ensureCapacity(text.length() + increase); while (end != StringUtils.INDEX_NOT_FOUND) { - _buffer.append(input, start, end).append(substitute); + _buffer.append(text, start, end).append(replacement); start = end + replLength; - if (--max == 0) { - break; - } - end = StringUtils.indexOf(input, searchString, start); + end = Strings.CS.indexOf(text, searchString, start); } - _buffer.append(input, start, input.length()); + _buffer.append(text, start, text.length()); return _buffer.toString(); } } From 4ff0448bb74fc9bdf91fd2db33015803a6885781 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Wed, 6 Aug 2025 19:43:07 -0700 Subject: [PATCH 094/167] fixing the bug in passive timeout handling (#16530) --- .../apache/pinot/query/service/dispatch/QueryDispatcher.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/QueryDispatcher.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/QueryDispatcher.java index 3cf132d69d77..2848054dee85 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/QueryDispatcher.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/service/dispatch/QueryDispatcher.java @@ -456,7 +456,7 @@ private static Map prepareRequestMetadata(long requestId, String requestMetadata.put(CommonConstants.Broker.Request.QueryOptionKey.TIMEOUT_MS, Long.toString(deadline.timeRemaining(TimeUnit.MILLISECONDS))); requestMetadata.put(CommonConstants.Broker.Request.QueryOptionKey.EXTRA_PASSIVE_TIMEOUT_MS, - Long.toString(QueryThreadContext.getPassiveDeadlineMs())); + Long.toString(QueryThreadContext.getPassiveDeadlineMs() - QueryThreadContext.getActiveDeadlineMs())); requestMetadata.putAll(queryOptions); return requestMetadata; } From 109d3917a89758416c653f0de785ed24226a78ac Mon Sep 17 00:00:00 2001 From: Rajat Venkatesh <1638298+vrajat@users.noreply.github.com> Date: Thu, 7 Aug 2025 13:00:28 +0530 Subject: [PATCH 095/167] Fix a flaky assert when testing OOM Cancellation of MSE Queries (#16533) --- .../tests/OfflineClusterMemBasedServerQueryKillingTest.java | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterMemBasedServerQueryKillingTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterMemBasedServerQueryKillingTest.java index efa34e69eb57..990a1b41cc68 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterMemBasedServerQueryKillingTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/OfflineClusterMemBasedServerQueryKillingTest.java @@ -291,7 +291,7 @@ public void testDigestOOMMSE() JsonNode queryResponse = postQuery(OOM_QUERY); String exceptionsNode = queryResponse.get("exceptions").toString(); assertTrue(exceptionsNode.contains("\"errorCode\":" + QueryErrorCode.INTERNAL.getId()), exceptionsNode); - assertTrue(exceptionsNode.contains("Received 1 error from stage 1"), exceptionsNode); + assertTrue(exceptionsNode.contains("Received 1 error"), exceptionsNode); assertTrue(_testAccountant.hasCallback()); assertEquals(queryResponse.get("requestId").asText(), _testAccountant.getQueryResourceTracker().getQueryId()); } @@ -327,7 +327,7 @@ public void testSelectionOnlyOOMMSE() String exceptionsNode = queryResponse.get("exceptions").toString(); assertTrue(exceptionsNode.contains("\"errorCode\":" + QueryErrorCode.INTERNAL.getId()), exceptionsNode); - assertTrue(exceptionsNode.contains("Received 1 error from stage 1"), exceptionsNode); + assertTrue(exceptionsNode.contains("Received 1 error"), exceptionsNode); assertTrue(_testAccountant.hasCallback()); assertEquals(queryResponse.get("requestId").asText(), _testAccountant.getQueryResourceTracker().getQueryId()); } @@ -347,7 +347,7 @@ public void testDigestOOM2MSE() JsonNode queryResponse = postQuery(OOM_QUERY_2); String exceptionsNode = queryResponse.get("exceptions").toString(); assertTrue(exceptionsNode.contains("\"errorCode\":" + QueryErrorCode.INTERNAL.getId()), exceptionsNode); - assertTrue(exceptionsNode.contains("Received 1 error from stage 1"), exceptionsNode); + assertTrue(exceptionsNode.contains("Received 1 error"), exceptionsNode); assertTrue(_testAccountant.hasCallback()); assertEquals(queryResponse.get("requestId").asText(), _testAccountant.getQueryResourceTracker().getQueryId()); } From b12d7f2f35143c2700d965b844926985baade48f Mon Sep 17 00:00:00 2001 From: NOOB <43700604+noob-se7en@users.noreply.github.com> Date: Thu, 7 Aug 2025 13:11:14 +0530 Subject: [PATCH 096/167] Allow and Initialise Server Rate Limiter before server starts consuming (#16407) * Enable server rate limiter during server startup * Renames method * Enable server rate limiter before server startup --- .../realtime/RealtimeConsumptionRateManager.java | 10 ++++------ .../server/starter/helix/BaseServerStarter.java | 14 +++++++------- 2 files changed, 11 insertions(+), 13 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeConsumptionRateManager.java b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeConsumptionRateManager.java index 95604153a688..7b747874ef98 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeConsumptionRateManager.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeConsumptionRateManager.java @@ -84,7 +84,7 @@ public static RealtimeConsumptionRateManager getInstance() { return InstanceHolder.INSTANCE; } - public void enableThrottling() { + public void enablePartitionRateLimiter() { _isThrottlingAllowed = true; } @@ -307,11 +307,9 @@ public ServerRateLimiter(double initialRateLimit, ServerMetrics serverMetrics, S } public void throttle(int numMsgsConsumed) { - if (InstanceHolder.INSTANCE._isThrottlingAllowed) { - _metricEmitter.record(numMsgsConsumed); // just incrementing counter (non-blocking) - if (numMsgsConsumed > 0) { - _rateLimiter.acquire(numMsgsConsumed); // blocks if needed - } + _metricEmitter.record(numMsgsConsumed); // just incrementing counter (non-blocking) + if (numMsgsConsumed > 0) { + _rateLimiter.acquire(numMsgsConsumed); // blocks if needed } } diff --git a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/BaseServerStarter.java b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/BaseServerStarter.java index 764a2c46a48f..979369a7e656 100644 --- a/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/BaseServerStarter.java +++ b/pinot-server/src/main/java/org/apache/pinot/server/starter/helix/BaseServerStarter.java @@ -686,6 +686,12 @@ public void start() InstanceDataManager instanceDataManager = _serverInstance.getInstanceDataManager(); instanceDataManager.setSupplierOfIsServerReadyToServeQueries(() -> _isServerReadyToServeQueries); + // Enable Server level realtime ingestion rate limier + RealtimeConsumptionRateManager.getInstance().createServerRateLimiter(_serverConf, serverMetrics); + PinotClusterConfigChangeListener serverRateLimitConfigChangeListener = + new ServerRateLimitConfigChangeListener(serverMetrics); + _clusterConfigChangeHandler.registerClusterConfigChangeListener(serverRateLimitConfigChangeListener); + initSegmentFetcher(_serverConf); StateModelFactory stateModelFactory = new SegmentOnlineOfflineStateModelFactory(_instanceId, instanceDataManager); @@ -777,12 +783,6 @@ public void start() preServeQueries(); - // Enable Server level realtime ingestion rate limier - RealtimeConsumptionRateManager.getInstance().createServerRateLimiter(_serverConf, serverMetrics); - PinotClusterConfigChangeListener serverRateLimitConfigChangeListener = - new ServerRateLimitConfigChangeListener(serverMetrics); - _clusterConfigChangeHandler.registerClusterConfigChangeListener(serverRateLimitConfigChangeListener); - // Start the thread accountant Tracing.ThreadAccountantOps.startThreadAccountant(); PinotClusterConfigChangeListener threadAccountantListener = @@ -799,7 +799,7 @@ public void start() Collections.singletonMap(Helix.IS_SHUTDOWN_IN_PROGRESS, Boolean.toString(false))); _isServerReadyToServeQueries = true; // Throttling for realtime consumption is disabled up to this point to allow maximum consumption during startup time - RealtimeConsumptionRateManager.getInstance().enableThrottling(); + RealtimeConsumptionRateManager.getInstance().enablePartitionRateLimiter(); LOGGER.info("Pinot server ready"); From 9f74fc7d14808e85515ae4ddf906f2d805b3a714 Mon Sep 17 00:00:00 2001 From: Yash Mayya Date: Thu, 7 Aug 2025 15:43:45 +0530 Subject: [PATCH 097/167] Disable Calcite's AggregateUnionAggregate planner rule by default (#16515) --- .../query/QueryPlannerRuleOptionsTest.java | 60 +++++++++++++++++++ .../test/resources/queries/SetOpPlans.json | 15 +++-- .../pinot/spi/utils/CommonConstants.java | 3 +- 3 files changed, 71 insertions(+), 7 deletions(-) diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryPlannerRuleOptionsTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryPlannerRuleOptionsTest.java index 5b0c00a0816f..0bcf5e118380 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryPlannerRuleOptionsTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/QueryPlannerRuleOptionsTest.java @@ -470,4 +470,64 @@ public void testEnableSortJoinCopy() { + " PinotLogicalTableScan(table=[[default, b]])\n"); //@formatter:on } + + @Test + public void testAggregateUnionAggregateDisabledByDefault() { + // Verify that the AggregateUnionAggregateRule is disabled by default + //@formatter:off + String query = "EXPLAIN PLAN FOR " + + "SELECT * FROM " + + "(SELECT DISTINCT col1 FROM a) " + + "UNION " + + "(SELECT DISTINCT col1 FROM b)"; + //@formatter:on + + String explain = _queryEnvironment.explainQuery(query, RANDOM_REQUEST_ID_GEN.nextLong()); + + // The aggregates above the table scans should not be merged into the one above the UNION ALL + assertEquals(explain, + "Execution Plan\n" + + "PinotLogicalAggregate(group=[{0}], aggType=[FINAL])\n" + + " PinotLogicalExchange(distribution=[hash[0]])\n" + + " PinotLogicalAggregate(group=[{0}], aggType=[LEAF])\n" + + " LogicalUnion(all=[true])\n" + + " PinotLogicalExchange(distribution=[hash[0]])\n" + + " PinotLogicalAggregate(group=[{0}], aggType=[FINAL])\n" + + " PinotLogicalExchange(distribution=[hash[0]])\n" + + " PinotLogicalAggregate(group=[{0}], aggType=[LEAF])\n" + + " PinotLogicalTableScan(table=[[default, a]])\n" + + " PinotLogicalExchange(distribution=[hash[0]])\n" + + " PinotLogicalAggregate(group=[{0}], aggType=[FINAL])\n" + + " PinotLogicalExchange(distribution=[hash[0]])\n" + + " PinotLogicalAggregate(group=[{0}], aggType=[LEAF])\n" + + " PinotLogicalTableScan(table=[[default, b]])\n"); + } + + @Test + public void testAggregateUnionAggregateEnabled() { + // Verify that the AggregateUnionAggregateRule is disabled by default + //@formatter:off + String query = "EXPLAIN PLAN FOR " + + "SELECT * FROM " + + "(SELECT DISTINCT col1 FROM a) " + + "UNION " + + "(SELECT DISTINCT col1 FROM b)"; + //@formatter:on + + String explain = explainQueryWithRuleEnabled(query, PlannerRuleNames.AGGREGATE_UNION_AGGREGATE); + + // There shouldn't be aggregates above the table scans since they should be merged into the one above the UNION ALL + assertEquals(explain, + "Execution Plan\n" + + "PinotLogicalAggregate(group=[{0}], aggType=[FINAL])\n" + + " PinotLogicalExchange(distribution=[hash[0]])\n" + + " PinotLogicalAggregate(group=[{0}], aggType=[LEAF])\n" + + " LogicalUnion(all=[true])\n" + + " PinotLogicalExchange(distribution=[hash[0]])\n" + + " LogicalProject(col1=[$0])\n" + + " PinotLogicalTableScan(table=[[default, a]])\n" + + " PinotLogicalExchange(distribution=[hash[0]])\n" + + " LogicalProject(col1=[$0])\n" + + " PinotLogicalTableScan(table=[[default, b]])\n"); + } } diff --git a/pinot-query-planner/src/test/resources/queries/SetOpPlans.json b/pinot-query-planner/src/test/resources/queries/SetOpPlans.json index 91d796cff63d..59eff321c573 100644 --- a/pinot-query-planner/src/test/resources/queries/SetOpPlans.json +++ b/pinot-query-planner/src/test/resources/queries/SetOpPlans.json @@ -46,13 +46,16 @@ "\n PinotLogicalAggregate(group=[{0, 1}], aggType=[LEAF])", "\n LogicalUnion(all=[true])", "\n PinotLogicalExchange(distribution=[hash[0, 1]])", - "\n LogicalUnion(all=[true])", + "\n PinotLogicalAggregate(group=[{0, 1}], aggType=[FINAL])", "\n PinotLogicalExchange(distribution=[hash[0, 1]])", - "\n LogicalProject(col1=[$0], col2=[$1])", - "\n PinotLogicalTableScan(table=[[default, a]])", - "\n PinotLogicalExchange(distribution=[hash[0, 1]])", - "\n LogicalProject(col1=[$0], col2=[$1])", - "\n PinotLogicalTableScan(table=[[default, b]])", + "\n PinotLogicalAggregate(group=[{0, 1}], aggType=[LEAF])", + "\n LogicalUnion(all=[true])", + "\n PinotLogicalExchange(distribution=[hash[0, 1]])", + "\n LogicalProject(col1=[$0], col2=[$1])", + "\n PinotLogicalTableScan(table=[[default, a]])", + "\n PinotLogicalExchange(distribution=[hash[0, 1]])", + "\n LogicalProject(col1=[$0], col2=[$1])", + "\n PinotLogicalTableScan(table=[[default, b]])", "\n PinotLogicalExchange(distribution=[hash[0, 1]])", "\n LogicalProject(col1=[$0], col2=[$1])", "\n PinotLogicalTableScan(table=[[default, c]])", diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java index 367c93704465..3ab38c56a03a 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java @@ -820,7 +820,8 @@ public static class PlannerRuleNames { public static final Set DEFAULT_DISABLED_RULES = Set.of( PlannerRuleNames.AGGREGATE_JOIN_TRANSPOSE_EXTENDED, PlannerRuleNames.SORT_JOIN_TRANSPOSE, - PlannerRuleNames.SORT_JOIN_COPY + PlannerRuleNames.SORT_JOIN_COPY, + PlannerRuleNames.AGGREGATE_UNION_AGGREGATE ); public static class FailureDetector { From 24c23465bb89b8d3d7ca394e789b62b40f58be84 Mon Sep 17 00:00:00 2001 From: Sonam Mandal Date: Thu, 7 Aug 2025 09:53:08 -0700 Subject: [PATCH 098/167] Table rebalance changes for handling peer-download enabled tables (#16341) * Change rebalance pre-check to warn for peer download enabled rather than pauseless * Don't allow downtime=true rebalance or minAvailableReplicas=0 when RF>1 if peer-download is enabled * Initial changes for adding PeerDownloadHandler * Add forceDowntime flag to allow downtime rebalance for peer-download enabled tables * Address review comments * Fail downtime rebalance if data loss scenarios are found, update scenario to check for COMMITTING state * Address review comments * Add exception handling for getNextAssignment * Only perform data loss check if nextAssignment is different from current assignment * Replace forceDownload flag with allowPeerDownloadDataLoss, and remove checks to disallow downtime rebalance for peer-download enabled * Make downtime for peer-download more efficient if allowPeerDownloadDataLoss = true * add test to test PeerDownloadTableDataLossRiskAssessor in table rebalance * remove generateDataLossRiskMessage and refactor hasDataLossRisk to assessDataLossRisk * Address review comments * Fix up generated error message to remove mention of force commit --------- Co-authored-by: J-HowHuang --- .../resources/PinotTableRestletResource.java | 8 + .../PinotLLCRealtimeSegmentManager.java | 8 +- .../rebalance/DefaultRebalancePreChecker.java | 33 +- .../helix/core/rebalance/RebalanceConfig.java | 42 +- .../helix/core/rebalance/TableRebalancer.java | 228 +++++++- .../rebalance/DataLossRiskAssessorTest.java | 523 ++++++++++++++++++ .../TableRebalancerClusterStatelessTest.java | 110 +++- .../core/rebalance/TableRebalancerTest.java | 159 ++++-- 8 files changed, 999 insertions(+), 112 deletions(-) create mode 100644 pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/DataLossRiskAssessorTest.java diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTableRestletResource.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTableRestletResource.java index 5a3516595957..1a60c39bd742 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTableRestletResource.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTableRestletResource.java @@ -716,6 +716,13 @@ public RebalanceResult rebalance( @DefaultValue("false") @QueryParam("bootstrap") boolean bootstrap, @ApiParam(value = "Whether to allow downtime for the rebalance") @DefaultValue("false") @QueryParam("downtime") boolean downtime, + @ApiParam(value = "This flag only applies to peer-download enabled tables undergoing downtime=true or " + + "minAvailableReplicas=0 rebalance (both of which can result in possible data loss scenarios). If enabled, " + + "this flag will allow the rebalance to continue even in cases where data loss scenarios have been " + + "detected, otherwise the rebalance will be failed and user action will be required to rebalance again. " + + "This flag should be used with caution and only used in scenarios where data loss is acceptable") + @DefaultValue("false") @QueryParam("allowPeerDownloadDataLoss") + boolean allowPeerDownloadDataLoss, @ApiParam(value = "For no-downtime rebalance, minimum number of replicas to keep alive during rebalance, or " + "maximum number of replicas allowed to be unavailable if value is negative") @DefaultValue("-1") @QueryParam("minAvailableReplicas") int minAvailableReplicas, @@ -779,6 +786,7 @@ public RebalanceResult rebalance( rebalanceConfig.setMinimizeDataMovement(minimizeDataMovement); rebalanceConfig.setBootstrap(bootstrap); rebalanceConfig.setDowntime(downtime); + rebalanceConfig.setAllowPeerDownloadDataLoss(allowPeerDownloadDataLoss); rebalanceConfig.setMinAvailableReplicas(minAvailableReplicas); rebalanceConfig.setLowDiskMode(lowDiskMode); rebalanceConfig.setBestEfforts(bestEfforts); diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java index a87a2ea008c3..099a49a10393 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java @@ -2552,7 +2552,13 @@ public void repairSegmentsInErrorStateForPauselessConsumption(TableConfig tableC } } - private boolean allowRepairOfErrorSegments(boolean repairErrorSegmentsForPartialUpsertOrDedup, + /** + * Whether to allow repairing the ERROR segment or not + * @param repairErrorSegmentsForPartialUpsertOrDedup API context flag, if true then always allow repair + * @param tableConfig tableConfig + * @return Returns true if repair is allowed for ERROR segments or not + */ + public boolean allowRepairOfErrorSegments(boolean repairErrorSegmentsForPartialUpsertOrDedup, TableConfig tableConfig) { if (repairErrorSegmentsForPartialUpsertOrDedup) { // If API context has repairErrorSegmentsForPartialUpsertOrDedup=true, allow repair. diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreChecker.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreChecker.java index c1c360d59e6f..fe782f4bd3de 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreChecker.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/DefaultRebalancePreChecker.java @@ -33,7 +33,6 @@ import org.apache.pinot.common.assignment.InstanceAssignmentConfigUtils; import org.apache.pinot.common.exception.InvalidConfigException; import org.apache.pinot.common.restlet.resources.DiskUsageInfo; -import org.apache.pinot.common.utils.PauselessConsumptionUtils; import org.apache.pinot.controller.helix.core.PinotHelixResourceManager; import org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignmentUtils; import org.apache.pinot.controller.util.TableMetadataReader; @@ -352,12 +351,16 @@ private RebalancePreCheckerResult checkRebalanceConfig(RebalanceConfig rebalance List segmentsToMove = SegmentAssignmentUtils.getSegmentsToMove(currentAssignment, targetAssignment); int numReplicas = Integer.MAX_VALUE; - if (rebalanceConfig.isDowntime() || PauselessConsumptionUtils.isPauselessEnabled(tableConfig)) { + String peerSegmentDownloadScheme = tableConfig.getValidationConfig().getPeerSegmentDownloadScheme(); + if (rebalanceConfig.isDowntime() || peerSegmentDownloadScheme != null) { for (String segment : segmentsToMove) { numReplicas = Math.min(targetAssignment.get(segment).size(), numReplicas); } } + // For non-peer download enabled tables, warn if downtime is enabled but numReplicas > 1. Should only use + // downtime=true for such tables if downtime is indeed acceptable whereas for numReplicas = 1, rebalance cannot + // be done without downtime if (rebalanceConfig.isDowntime()) { if (!segmentsToMove.isEmpty() && numReplicas > 1) { pass = false; @@ -365,22 +368,32 @@ private RebalancePreCheckerResult checkRebalanceConfig(RebalanceConfig rebalance } } - // It was revealed a risk of data loss for pauseless tables during rebalance, when downtime=true or - // minAvailableReplicas=0 -- If a segment is being moved and has not yet uploaded to deep store, premature - // deletion could cause irrecoverable data loss. This pre-check added as a workaround to warn the potential risk. - // TODO: Get to the root cause of the issue and revisit this pre-check. - if (PauselessConsumptionUtils.isPauselessEnabled(tableConfig)) { + // Peer download enabled tables may have data loss during rebalance, when downtime=true or minAvailableReplicas=0. + // The scenario plays out as follows: + // 1. If the newly built consuming segment cannot be uploaded to deep store, it may set up the download URI + // as an empty string: "" + // 2. When this happens, other servers expect to download the segment from a peer server that built the segment or + // has a copy of the segment + // 3. With downtime rebalance (or if minAvailableReplicas=0), the IS may be updated for all the servers of a given + // segment + // 4. The above may lead to dropping the existing segments from the existing servers without waiting for the newly + // added servers to download the segment from the peer. In this case since a deep store copy does not exist, + // there is no way to recover this segment without manually re-building it + // Thus, to avoid the above data loss scenario, it is not recommended to run downtime rebalance for peer download + // enabled tables. This pre-check is added to warn of the potential risk. + if (peerSegmentDownloadScheme != null) { int minAvailableReplica = rebalanceConfig.getMinAvailableReplicas(); if (minAvailableReplica < 0) { minAvailableReplica = numReplicas + minAvailableReplica; } if (numReplicas == 1) { pass = false; - warnings.add("Replication of the table is 1, which is not recommended for pauseless tables as it may cause " - + "data loss during rebalance"); + warnings.add("Replication of the table is 1, which is not recommended for peer-download enabled tables as it " + + "may cause data loss during rebalance"); } else if (rebalanceConfig.isDowntime() || minAvailableReplica <= 0) { pass = false; - warnings.add("Downtime or minAvailableReplicas=0 for pauseless tables may cause data loss during rebalance"); + warnings.add("Downtime or minAvailableReplicas<=0 for peer-download enabled tables may cause data loss during " + + "rebalance"); } } diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceConfig.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceConfig.java index 38f1fff33d08..de6cdb7ea07a 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceConfig.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/RebalanceConfig.java @@ -65,6 +65,15 @@ public class RebalanceConfig { @ApiModelProperty(example = "false") private boolean _downtime = false; + // This flag only applies to peer-download enabled tables undergoing downtime=true or minAvailableReplicas=0 + // rebalance (both of which can result in possible data loss scenarios). If enabled, this flag will allow the + // rebalance to continue even in cases where data loss scenarios have been detected, otherwise the rebalance will + // be failed and user action will be required to rebalance again. This flag should be used with caution and only + // used in scenarios where data loss is acceptable. + @JsonProperty("allowPeerDownloadDataLoss") + @ApiModelProperty(example = "false") + private boolean _allowPeerDownloadDataLoss = false; + // For no-downtime rebalance, minimum number of replicas to keep alive during rebalance, or maximum number of replicas // allowed to be unavailable if value is negative @JsonProperty("minAvailableReplicas") @@ -209,6 +218,14 @@ public void setDowntime(boolean downtime) { _downtime = downtime; } + public boolean isAllowPeerDownloadDataLoss() { + return _allowPeerDownloadDataLoss; + } + + public void setAllowPeerDownloadDataLoss(boolean allowPeerDownloadDataLoss) { + _allowPeerDownloadDataLoss = allowPeerDownloadDataLoss; + } + public int getMinAvailableReplicas() { return _minAvailableReplicas; } @@ -351,23 +368,25 @@ public void setDiskUtilizationThreshold(double diskUtilizationThreshold) { public String toString() { return "RebalanceConfig{" + "_dryRun=" + _dryRun + ", preChecks=" + _preChecks + ", _reassignInstances=" + _reassignInstances + ", _includeConsuming=" + _includeConsuming + ", _minimizeDataMovement=" - + _minimizeDataMovement + ", _bootstrap=" + _bootstrap + ", _downtime=" + _downtime + ", _minAvailableReplicas=" - + _minAvailableReplicas + ", _bestEfforts=" + _bestEfforts + ", batchSizePerServer=" + _batchSizePerServer - + ", _externalViewCheckIntervalInMs=" + _externalViewCheckIntervalInMs - + ", _externalViewStabilizationTimeoutInMs=" + _externalViewStabilizationTimeoutInMs - + ", _updateTargetTier=" + _updateTargetTier + ", _heartbeatIntervalInMs=" + _heartbeatIntervalInMs - + ", _heartbeatTimeoutInMs=" + _heartbeatTimeoutInMs + ", _maxAttempts=" + _maxAttempts - + ", _retryInitialDelayInMs=" + _retryInitialDelayInMs + ", _diskUtilizationThreshold=" - + _diskUtilizationThreshold + ", _forceCommit=" + _forceCommit + ", _forceCommitBatchSize=" - + _forceCommitBatchSize + ", _forceCommitBatchStatusCheckIntervalMs=" + _forceCommitBatchStatusCheckIntervalMs + + _minimizeDataMovement + ", _bootstrap=" + _bootstrap + ", _downtime=" + _downtime + + ", _allowPeerDownloadDataLoss=" + _allowPeerDownloadDataLoss + ", _minAvailableReplicas=" + + _minAvailableReplicas + ", _bestEfforts=" + _bestEfforts + ", batchSizePerServer=" + + _batchSizePerServer + ", _externalViewCheckIntervalInMs=" + _externalViewCheckIntervalInMs + + ", _externalViewStabilizationTimeoutInMs=" + _externalViewStabilizationTimeoutInMs + ", _updateTargetTier=" + + _updateTargetTier + ", _heartbeatIntervalInMs=" + _heartbeatIntervalInMs + ", _heartbeatTimeoutInMs=" + + _heartbeatTimeoutInMs + ", _maxAttempts=" + _maxAttempts + ", _retryInitialDelayInMs=" + + _retryInitialDelayInMs + ", _diskUtilizationThreshold=" + _diskUtilizationThreshold + ", _forceCommit=" + + _forceCommit + ", _forceCommitBatchSize=" + _forceCommitBatchSize + + ", _forceCommitBatchStatusCheckIntervalMs=" + _forceCommitBatchStatusCheckIntervalMs + ", _forceCommitBatchStatusCheckTimeoutMs=" + _forceCommitBatchStatusCheckTimeoutMs + '}'; } public String toQueryString() { return "dryRun=" + _dryRun + "&preChecks=" + _preChecks + "&reassignInstances=" + _reassignInstances + "&includeConsuming=" + _includeConsuming + "&bootstrap=" + _bootstrap + "&downtime=" + _downtime - + "&minAvailableReplicas=" + _minAvailableReplicas + "&bestEfforts=" + _bestEfforts - + "&minimizeDataMovement=" + _minimizeDataMovement.name() + "&batchSizePerServer=" + _batchSizePerServer + + "&allowPeerDownloadDataLoss=" + _allowPeerDownloadDataLoss + "&minAvailableReplicas=" + _minAvailableReplicas + + "&bestEfforts=" + _bestEfforts + "&minimizeDataMovement=" + _minimizeDataMovement.name() + + "&batchSizePerServer=" + _batchSizePerServer + "&externalViewCheckIntervalInMs=" + _externalViewCheckIntervalInMs + "&externalViewStabilizationTimeoutInMs=" + _externalViewStabilizationTimeoutInMs + "&updateTargetTier=" + _updateTargetTier + "&heartbeatIntervalInMs=" + _heartbeatIntervalInMs @@ -387,6 +406,7 @@ public static RebalanceConfig copy(RebalanceConfig cfg) { rc._includeConsuming = cfg._includeConsuming; rc._bootstrap = cfg._bootstrap; rc._downtime = cfg._downtime; + rc._allowPeerDownloadDataLoss = cfg._allowPeerDownloadDataLoss; rc._minAvailableReplicas = cfg._minAvailableReplicas; rc._bestEfforts = cfg._bestEfforts; rc._minimizeDataMovement = cfg._minimizeDataMovement; diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java index d238a54d967a..24854b8bb75f 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancer.java @@ -66,6 +66,7 @@ import org.apache.pinot.common.tier.PinotServerTierStorage; import org.apache.pinot.common.tier.Tier; import org.apache.pinot.common.tier.TierFactory; +import org.apache.pinot.common.utils.PauselessConsumptionUtils; import org.apache.pinot.common.utils.SegmentUtils; import org.apache.pinot.common.utils.config.TierConfigUtils; import org.apache.pinot.controller.api.resources.ForceCommitBatchConfig; @@ -90,6 +91,7 @@ import org.apache.pinot.spi.stream.StreamConsumerFactoryProvider; import org.apache.pinot.spi.stream.StreamMetadataProvider; import org.apache.pinot.spi.stream.StreamPartitionMsgOffset; +import org.apache.pinot.spi.utils.CommonConstants; import org.apache.pinot.spi.utils.CommonConstants.Helix.StateModel.SegmentStateModel; import org.apache.pinot.spi.utils.Enablement; import org.apache.pinot.spi.utils.IngestionConfigUtils; @@ -223,6 +225,7 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb boolean includeConsuming = rebalanceConfig.isIncludeConsuming(); boolean bootstrap = rebalanceConfig.isBootstrap(); boolean downtime = rebalanceConfig.isDowntime(); + boolean allowPeerDownloadDataLoss = rebalanceConfig.isAllowPeerDownloadDataLoss(); int minReplicasToKeepUpForNoDowntime = rebalanceConfig.getMinAvailableReplicas(); boolean lowDiskMode = rebalanceConfig.isLowDiskMode(); boolean bestEfforts = rebalanceConfig.isBestEfforts(); @@ -242,12 +245,12 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb } tableRebalanceLogger.info( "Start rebalancing with dryRun: {}, preChecks: {}, reassignInstances: {}, " - + "includeConsuming: {}, bootstrap: {}, downtime: {}, minReplicasToKeepUpForNoDowntime: {}, " - + "enableStrictReplicaGroup: {}, lowDiskMode: {}, bestEfforts: {}, batchSizePerServer: {}, " - + "externalViewCheckIntervalInMs: {}, externalViewStabilizationTimeoutInMs: {}, minimizeDataMovement: {}, " - + "forceCommit: {}, forceCommitBatchSize: {}, forceCommitBatchStatusCheckIntervalMs: {}, " - + "forceCommitBatchStatusCheckTimeoutMs: {}", - dryRun, preChecks, reassignInstances, includeConsuming, bootstrap, downtime, + + "includeConsuming: {}, bootstrap: {}, downtime: {}, allowPeerDownloadDataLoss: {}, " + + "minReplicasToKeepUpForNoDowntime: {}, enableStrictReplicaGroup: {}, lowDiskMode: {}, bestEfforts: {}, " + + "batchSizePerServer: {}, externalViewCheckIntervalInMs: {}, externalViewStabilizationTimeoutInMs: {}, " + + "minimizeDataMovement: {}, forceCommit: {}, forceCommitBatchSize: {}, " + + "forceCommitBatchStatusCheckIntervalMs: {}, forceCommitBatchStatusCheckTimeoutMs: {}", + dryRun, preChecks, reassignInstances, includeConsuming, bootstrap, downtime, allowPeerDownloadDataLoss, minReplicasToKeepUpForNoDowntime, enableStrictReplicaGroup, lowDiskMode, bestEfforts, batchSizePerServer, externalViewCheckIntervalInMs, externalViewStabilizationTimeoutInMs, minimizeDataMovement, forceCommit, rebalanceConfig.getForceCommitBatchSize(), @@ -387,6 +390,8 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb tierToInstancePartitionsMap, targetAssignment, preChecksResult, summaryResult); } + String peerSegmentDownloadScheme = tableConfig.getValidationConfig().getPeerSegmentDownloadScheme(); + if (downtime) { tableRebalanceLogger.info("Rebalancing with downtime"); if (forceCommit) { @@ -411,6 +416,29 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb tierToInstancePartitionsMap, rebalanceConfig); } } + + // If peer-download is enabled, verify that for all segments with changes in assignment, it is safe to rebalance + // Create the DataLossRiskAssessor which is used to check for data loss scenarios if peer-download is enabled + // for a table. Skip this step if allowPeerDownloadDataLoss = true + if (peerSegmentDownloadScheme != null && !allowPeerDownloadDataLoss) { + DataLossRiskAssessor dataLossRiskAssessor = new PeerDownloadTableDataLossRiskAssessor(tableNameWithType, + tableConfig, _helixManager, _pinotLLCRealtimeSegmentManager); + for (Map.Entry> segmentToAssignment : currentAssignment.entrySet()) { + String segmentName = segmentToAssignment.getKey(); + Map assignment = segmentToAssignment.getValue(); + if (!assignment.equals(targetAssignment.get(segmentName))) { + Pair dataLossResult = dataLossRiskAssessor.assessDataLossRisk(segmentName); + if (dataLossResult.getLeft()) { + // Fail the rebalance if a segment with the potential for data loss is found + String errorMsg = dataLossResult.getRight(); + onReturnFailure(errorMsg, new IllegalStateException(errorMsg), tableRebalanceLogger); + return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.FAILED, errorMsg, instancePartitionsMap, + tierToInstancePartitionsMap, targetAssignment, preChecksResult, summaryResult); + } + } + } + } + // Reuse current IdealState to update the IdealState in cluster ZNRecord idealStateRecord = currentIdealState.getRecord(); idealStateRecord.setMapFields(targetAssignment); @@ -489,6 +517,19 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb minAvailableReplicas = numCurrentAssignmentReplicas; } + DataLossRiskAssessor dataLossRiskAssessor; + if (minAvailableReplicas == 0 && peerSegmentDownloadScheme != null && !allowPeerDownloadDataLoss) { + // Create the DataLossRiskAssessor which is used to check for data loss scenarios if peer-download is enabled + // for a table + dataLossRiskAssessor = new PeerDownloadTableDataLossRiskAssessor(tableNameWithType, tableConfig, _helixManager, + _pinotLLCRealtimeSegmentManager); + } else { + // If peer-download is disabled or minAvailableReplicas > 0, there is no data loss risk so create a no-op + // assessor. If allowPeerDownloadDataLoss = true, then also skip checking for data loss since the caller's + // intent is to rebalance in spite of data loss + dataLossRiskAssessor = new NoOpRiskAssessor(); + } + tableRebalanceLogger.info("Rebalancing with minAvailableReplicas: {}, enableStrictReplicaGroup: {}, " + "bestEfforts: {}, externalViewCheckIntervalInMs: {}, externalViewStabilizationTimeoutInMs: {}", minAvailableReplicas, enableStrictReplicaGroup, bestEfforts, externalViewCheckIntervalInMs, @@ -503,6 +544,7 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb PartitionIdFetcher partitionIdFetcher = new PartitionIdFetcherImpl(tableNameWithType, TableConfigUtils.getPartitionColumn(tableConfig), _helixManager, isStrictRealtimeSegmentAssignment); + // We repeat the following steps until the target assignment is reached: // 1. Wait for ExternalView to converge with the IdealState. Fail the rebalance if it doesn't make progress within // the timeout. @@ -613,9 +655,18 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb // Step 2: Handle force commit if flag is set, then recalculate if force commit occurred if (shouldForceCommit) { - nextAssignment = - getNextAssignment(currentAssignment, targetAssignment, minAvailableReplicas, enableStrictReplicaGroup, - lowDiskMode, batchSizePerServer, segmentPartitionIdMap, partitionIdFetcher, tableRebalanceLogger); + try { + nextAssignment = + getNextAssignment(currentAssignment, targetAssignment, minAvailableReplicas, enableStrictReplicaGroup, + lowDiskMode, batchSizePerServer, segmentPartitionIdMap, partitionIdFetcher, dataLossRiskAssessor, + tableRebalanceLogger); + } catch (Exception e) { + String errorMsg = + "Caught exception while calculating the next assignment, aborting the rebalance: " + e.getMessage(); + onReturnFailure(errorMsg, e, tableRebalanceLogger); + return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.FAILED, errorMsg, instancePartitionsMap, + tierToInstancePartitionsMap, targetAssignment, preChecksResult, summaryResult); + } Set consumingSegmentsToMoveNext = getMovingConsumingSegments(currentAssignment, nextAssignment); if (!consumingSegmentsToMoveNext.isEmpty()) { @@ -669,9 +720,18 @@ private RebalanceResult doRebalance(TableConfig tableConfig, RebalanceConfig reb tierToInstancePartitionsMap, targetAssignment, preChecksResult, summaryResult); } - nextAssignment = - getNextAssignment(currentAssignment, targetAssignment, minAvailableReplicas, enableStrictReplicaGroup, - lowDiskMode, batchSizePerServer, segmentPartitionIdMap, partitionIdFetcher, tableRebalanceLogger); + try { + nextAssignment = + getNextAssignment(currentAssignment, targetAssignment, minAvailableReplicas, enableStrictReplicaGroup, + lowDiskMode, batchSizePerServer, segmentPartitionIdMap, partitionIdFetcher, dataLossRiskAssessor, + tableRebalanceLogger); + } catch (Exception e) { + String errorMsg = + "Caught exception while calculating the next assignment, aborting the rebalance: " + e.getMessage(); + onReturnFailure(errorMsg, e, tableRebalanceLogger); + return new RebalanceResult(rebalanceJobId, RebalanceResult.Status.FAILED, errorMsg, instancePartitionsMap, + tierToInstancePartitionsMap, targetAssignment, preChecksResult, summaryResult); + } tableRebalanceLogger.info( "Got the next assignment with number of segments to be added/removed for each instance: {}", SegmentAssignmentUtils.getNumSegmentsToMovePerInstance(currentAssignment, nextAssignment)); @@ -1141,7 +1201,7 @@ public Pair, Boolean> getInstanc getInstancePartitions(tableConfig, InstancePartitionsType.COMPLETED, reassignInstances, bootstrap, dryRun, minimizeDataMovement, tableRebalanceLogger); tableRebalanceLogger.info("COMPLETED segments should be relocated, fetching/computing COMPLETED instance " - + "partitions for table: {}", tableNameWithType); + + "partitions for table: {}", tableNameWithType); instancePartitionsMap.put(InstancePartitionsType.COMPLETED, partitionAndUnchangedForCompleted.getLeft()); instancePartitionsUnchanged &= partitionAndUnchangedForCompleted.getRight(); } else { @@ -1227,7 +1287,7 @@ private Pair getInstancePartitions(TableConfig tabl return Pair.of(instancePartitions, instancePartitionsUnchanged); } else { tableRebalanceLogger.info("{} instance assignment is not allowed, using default instance partitions for " - + "table: {}", instancePartitionsType, tableNameWithType); + + "table: {}", instancePartitionsType, tableNameWithType); InstancePartitions instancePartitions = InstancePartitionsUtils.computeDefaultInstancePartitions(_helixManager, tableConfig, instancePartitionsType); @@ -1380,9 +1440,8 @@ private IdealState waitForExternalViewToConverge(String tableNameWithType, boole // Update unique segment list as IS-EV trigger must have processed these allSegmentsFromIdealState = idealState.getRecord().getMapFields().keySet(); if (_tableRebalanceObserver.isStopped()) { - throw new RuntimeException( - String.format("Rebalance has already stopped with status: %s", - _tableRebalanceObserver.getStopStatus())); + throw new RuntimeException(String.format("Rebalance has already stopped with status: %s", + _tableRebalanceObserver.getStopStatus())); } if (isExternalViewConverged(externalView.getRecord().getMapFields(), idealState.getRecord().getMapFields(), lowDiskMode, bestEfforts, segmentsToMonitor, tableRebalanceLogger)) { @@ -1574,9 +1633,9 @@ private static void handleErrorInstance(String segmentName, String instanceName, static Map> getNextAssignment(Map> currentAssignment, Map> targetAssignment, int minAvailableReplicas, boolean enableStrictReplicaGroup, boolean lowDiskMode, int batchSizePerServer, Object2IntOpenHashMap segmentPartitionIdMap, - PartitionIdFetcher partitionIdFetcher) { + PartitionIdFetcher partitionIdFetcher, DataLossRiskAssessor dataLossRiskAssessor) { return getNextAssignment(currentAssignment, targetAssignment, minAvailableReplicas, enableStrictReplicaGroup, - lowDiskMode, batchSizePerServer, segmentPartitionIdMap, partitionIdFetcher, LOGGER); + lowDiskMode, batchSizePerServer, segmentPartitionIdMap, partitionIdFetcher, dataLossRiskAssessor, LOGGER); } /** @@ -1602,19 +1661,19 @@ static Map> getNextAssignment(Map> getNextAssignment(Map> currentAssignment, Map> targetAssignment, int minAvailableReplicas, boolean enableStrictReplicaGroup, boolean lowDiskMode, int batchSizePerServer, Object2IntOpenHashMap segmentPartitionIdMap, - PartitionIdFetcher partitionIdFetcher, Logger tableRebalanceLogger) { + PartitionIdFetcher partitionIdFetcher, DataLossRiskAssessor dataLossRiskAssessor, Logger tableRebalanceLogger) { return enableStrictReplicaGroup ? getNextStrictReplicaGroupAssignment(currentAssignment, targetAssignment, minAvailableReplicas, lowDiskMode, - batchSizePerServer, segmentPartitionIdMap, partitionIdFetcher, tableRebalanceLogger) + batchSizePerServer, segmentPartitionIdMap, partitionIdFetcher, dataLossRiskAssessor, tableRebalanceLogger) : getNextNonStrictReplicaGroupAssignment(currentAssignment, targetAssignment, minAvailableReplicas, - lowDiskMode, batchSizePerServer); + lowDiskMode, batchSizePerServer, dataLossRiskAssessor); } private static Map> getNextStrictReplicaGroupAssignment( Map> currentAssignment, Map> targetAssignment, int minAvailableReplicas, boolean lowDiskMode, int batchSizePerServer, Object2IntOpenHashMap segmentPartitionIdMap, PartitionIdFetcher partitionIdFetcher, - Logger tableRebalanceLogger) { + DataLossRiskAssessor dataLossRiskAssessor, Logger tableRebalanceLogger) { Map> nextAssignment = new TreeMap<>(); Map numSegmentsToOffloadMap = getNumSegmentsToOffloadMap(currentAssignment, targetAssignment); Map, Set>, Set> assignmentMap = new HashMap<>(); @@ -1625,7 +1684,7 @@ private static Map> getNextStrictReplicaGroupAssignm // Directly update the nextAssignment with anyServerExhaustedBatchSize = false and return if batching is disabled updateNextAssignmentForPartitionIdStrictReplicaGroup(currentAssignment, targetAssignment, nextAssignment, false, minAvailableReplicas, lowDiskMode, numSegmentsToOffloadMap, assignmentMap, - availableInstancesMap, serverToNumSegmentsAddedSoFar); + availableInstancesMap, serverToNumSegmentsAddedSoFar, dataLossRiskAssessor); return nextAssignment; } @@ -1670,7 +1729,7 @@ private static Map> getNextStrictReplicaGroupAssignm } updateNextAssignmentForPartitionIdStrictReplicaGroup(curAssignment, targetAssignment, nextAssignment, anyServerExhaustedBatchSize, minAvailableReplicas, lowDiskMode, numSegmentsToOffloadMap, assignmentMap, - availableInstancesMap, serverToNumSegmentsAddedSoFar); + availableInstancesMap, serverToNumSegmentsAddedSoFar, dataLossRiskAssessor); } } @@ -1684,7 +1743,8 @@ private static void updateNextAssignmentForPartitionIdStrictReplicaGroup( Map> nextAssignment, boolean anyServerExhaustedBatchSize, int minAvailableReplicas, boolean lowDiskMode, Map numSegmentsToOffloadMap, Map, Set>, Set> assignmentMap, - Map, Set> availableInstancesMap, Map serverToNumSegmentsAddedSoFar) { + Map, Set> availableInstancesMap, Map serverToNumSegmentsAddedSoFar, + DataLossRiskAssessor dataLossRiskAssessor) { if (anyServerExhaustedBatchSize) { // Exhausted the batch size for at least 1 server, just copy over the remaining segments as is nextAssignment.putAll(currentAssignment); @@ -1728,6 +1788,13 @@ private static void updateNextAssignmentForPartitionIdStrictReplicaGroup( Set serversAddedForSegment = getServersAddedInSingleSegmentAssignment(currentInstanceStateMap, nextAssignment.get(segmentName)); serversAddedForSegment.forEach(server -> serverToNumSegmentsAddedSoFar.merge(server, 1, Integer::sum)); + + // Since next assignment doesn't match current assignment, it means the segment will be moved. Check if there + // is a data loss risk + Pair dataLossRiskResult = dataLossRiskAssessor.assessDataLossRisk(segmentName); + if (dataLossRiskResult.getLeft()) { + throw new IllegalStateException(dataLossRiskResult.getRight()); + } } } } @@ -1813,9 +1880,107 @@ public int fetch(String segmentName) { } } + @VisibleForTesting + @FunctionalInterface + interface DataLossRiskAssessor { + Pair NO_DATA_LOSS_RISK_RESULT = Pair.of(false, null); + + /** + * Assess the risk of data loss for the given segment. + * + * @param segmentName Name of the segment to assess + * @return A pair where the first element indicates if there is a risk of data loss, and the second element is a + * message describing the risk (if any). + */ + Pair assessDataLossRisk(String segmentName); + } + + /** + * To be used for non-peer download enabled tables or peer-download enabled tables rebalanced with + * minAvailableReplicas > 0 + */ + @VisibleForTesting + static class NoOpRiskAssessor implements DataLossRiskAssessor { + NoOpRiskAssessor() { + } + + @Override + public Pair assessDataLossRisk(String segmentName) { + return NO_DATA_LOSS_RISK_RESULT; + } + } + + /** + * To be used for peer-download enabled tables with downtime=true or minAvailableReplicas=0 + */ + @VisibleForTesting + static class PeerDownloadTableDataLossRiskAssessor implements DataLossRiskAssessor { + private final String _tableNameWithType; + private final TableConfig _tableConfig; + private final HelixManager _helixManager; + private final PinotLLCRealtimeSegmentManager _pinotLLCRealtimeSegmentManager; + private final boolean _isPauselessEnabled; + + @VisibleForTesting + PeerDownloadTableDataLossRiskAssessor(String tableNameWithType, TableConfig tableConfig, HelixManager helixManager, + PinotLLCRealtimeSegmentManager pinotLLCRealtimeSegmentManager) { + // Should only be created for peer-download enabled tables with minAvailableReplicas = 0 + Preconditions.checkState(tableConfig.getValidationConfig().getPeerSegmentDownloadScheme() != null); + _tableNameWithType = tableNameWithType; + _tableConfig = tableConfig; + _helixManager = helixManager; + _pinotLLCRealtimeSegmentManager = pinotLLCRealtimeSegmentManager; + _isPauselessEnabled = PauselessConsumptionUtils.isPauselessEnabled(tableConfig); + } + + @Override + public Pair assessDataLossRisk(String segmentName) { + SegmentZKMetadata segmentZKMetadata = ZKMetadataProvider + .getSegmentZKMetadata(_helixManager.getHelixPropertyStore(), _tableNameWithType, segmentName); + if (segmentZKMetadata == null) { + return NO_DATA_LOSS_RISK_RESULT; + } + + // If the segment state is COMPLETED and the download URL is empty, there is a data loss risk + if (segmentZKMetadata.getStatus().isCompleted() && CommonConstants.Segment.METADATA_URI_FOR_PEER_DOWNLOAD.equals( + segmentZKMetadata.getDownloadUrl())) { + return Pair.of(true, generateDataLossRiskMessage(segmentName, true)); + } + + // If the segment is not yet completed, then the following scenarios are possible: + // - Non-upsert / non-dedup table: + // - data loss scenarios are not possible. Either the segment will restart consumption or the + // RealtimeSegmentValidationManager will kick in to fix up the segment if pauseless is enabled + // - Upsert / dedup table: + // - For non-pauseless tables, it is safe to move the segment without data loss concerns + // - For pauseless tables, if the segment is still in CONSUMING state, moving it is safe, but if it is in + // COMMITTING state then there is a risk of data loss on segment build failures as well since the + // RealtimeSegmentValidationManager does not automatically try to fix up these segments. To be safe it is + // best to return that there is a risk of data loss for pauseless enabled tables for segments in COMMITTING + // state + if (_isPauselessEnabled && segmentZKMetadata.getStatus() == CommonConstants.Segment.Realtime.Status.COMMITTING + && !_pinotLLCRealtimeSegmentManager.allowRepairOfErrorSegments(false, _tableConfig)) { + return Pair.of(true, generateDataLossRiskMessage(segmentName, false)); + } + return NO_DATA_LOSS_RISK_RESULT; + } + + private static String generateDataLossRiskMessage(String segmentName, boolean isCompletedSegment) { + if (isCompletedSegment) { + return "Moving segment " + segmentName + " as part of rebalance is risky for peer-download enabled tables, " + + "as the download URL is empty. Ensure that the deep store has a copy of the segment. It is recommended " + + "to pause ingestion prior to trying to rebalance such tables with downtime"; + } + return "Moving segment " + segmentName + " as part of rebalance is risky for peer-download enabled tables " + + "as it is in COMMITING state and repair is not allowed (it may be an upsert / dedup enabled table). It " + + "is recommended to pause ingestion prior to trying to rebalance such tables with downtime"; + } + } + private static Map> getNextNonStrictReplicaGroupAssignment( Map> currentAssignment, Map> targetAssignment, - int minAvailableReplicas, boolean lowDiskMode, int batchSizePerServer) { + int minAvailableReplicas, boolean lowDiskMode, int batchSizePerServer, + DataLossRiskAssessor dataLossRiskAssessor) { Map serverToNumSegmentsAddedSoFar = new HashMap<>(); Map> nextAssignment = new TreeMap<>(); Map numSegmentsToOffloadMap = getNumSegmentsToOffloadMap(currentAssignment, targetAssignment); @@ -1847,6 +2012,15 @@ private static Map> getNextNonStrictReplicaGroupAssi nextAssignment.put(segmentName, nextInstanceStateMap); updateNumSegmentsToOffloadMap(numSegmentsToOffloadMap, currentInstanceStateMap.keySet(), nextInstanceStateMap.keySet()); + + if (!nextAssignment.get(segmentName).equals(currentInstanceStateMap)) { + // Since next assignment doesn't match current assignment, it means the segment will be moved. Check if there + // is a data loss risk + Pair dataLossRiskResult = dataLossRiskAssessor.assessDataLossRisk(segmentName); + if (dataLossRiskResult.getLeft()) { + throw new IllegalStateException(dataLossRiskResult.getRight()); + } + } } } return nextAssignment; diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/DataLossRiskAssessorTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/DataLossRiskAssessorTest.java new file mode 100644 index 000000000000..1d76d1ad645a --- /dev/null +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/DataLossRiskAssessorTest.java @@ -0,0 +1,523 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.pinot.controller.helix.core.rebalance; + +import java.util.Collections; +import org.apache.commons.lang3.tuple.Pair; +import org.apache.helix.HelixManager; +import org.apache.helix.store.zk.ZkHelixPropertyStore; +import org.apache.helix.zookeeper.datamodel.ZNRecord; +import org.apache.pinot.common.metadata.ZKMetadataProvider; +import org.apache.pinot.common.metadata.segment.SegmentZKMetadata; +import org.apache.pinot.controller.helix.ControllerTest; +import org.apache.pinot.core.realtime.impl.fakestream.FakeStreamConfigUtils; +import org.apache.pinot.spi.config.table.DedupConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.config.table.UpsertConfig; +import org.apache.pinot.spi.config.table.ingestion.IngestionConfig; +import org.apache.pinot.spi.config.table.ingestion.StreamIngestionConfig; +import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.apache.pinot.spi.utils.builder.TableNameBuilder; +import org.mockito.MockedStatic; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertThrows; +import static org.testng.Assert.assertTrue; + + +public class DataLossRiskAssessorTest extends ControllerTest { + private static final String RAW_TABLE_NAME = "testTable"; + private static final String REALTIME_TABLE_NAME = TableNameBuilder.REALTIME.tableNameWithType(RAW_TABLE_NAME); + private static final int NUM_REPLICAS = 3; + + @BeforeClass + public void setUp() + throws Exception { + startZk(); + startController(getDefaultControllerConfiguration()); + addFakeBrokerInstancesToAutoJoinHelixCluster(1, true); + } + + @AfterClass + public void tearDown() { + stopFakeInstances(); + stopController(); + stopZk(); + } + + @Test + public void testDataLossRiskAssessorPeerDownloadDisabled() { + TableConfig tableConfig = + new TableConfigBuilder(TableType.REALTIME).setTableName(RAW_TABLE_NAME).setNumReplicas(NUM_REPLICAS).build(); + assertThrows(IllegalStateException.class, + () -> new TableRebalancer.PeerDownloadTableDataLossRiskAssessor(REALTIME_TABLE_NAME, tableConfig, + _helixManager, _helixResourceManager.getRealtimeSegmentManager())); + + assertThrows(IllegalStateException.class, + () -> new TableRebalancer.PeerDownloadTableDataLossRiskAssessor(REALTIME_TABLE_NAME, tableConfig, + _helixManager, _helixResourceManager.getRealtimeSegmentManager())); + + TableRebalancer.NoOpRiskAssessor noOpRiskAssessor = new TableRebalancer.NoOpRiskAssessor(); + Pair dataLossRiskResult = noOpRiskAssessor.assessDataLossRisk("randomSegmentName"); + assertFalse(dataLossRiskResult.getLeft()); + assertNull(dataLossRiskResult.getRight()); + } + + @Test + public void testDataLossRiskAssessorPeerDownloadEnabledCompletedSegment() { + String segmentName = "randomSegmentName"; + TableConfig tableConfig = + new TableConfigBuilder(TableType.REALTIME).setTableName(RAW_TABLE_NAME).setNumReplicas(NUM_REPLICAS).build(); + tableConfig.getValidationConfig().setPeerSegmentDownloadScheme("http"); + + try (MockedStatic zkMetadataProviderMockedStatic = mockStatic(ZKMetadataProvider.class)) { + // Create segment with non-empty download URL + SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segmentName); + segmentZKMetadata.setStatus(CommonConstants.Segment.Realtime.Status.DONE); + segmentZKMetadata.setDownloadUrl("nonEmptyDownloadURL"); + + zkMetadataProviderMockedStatic.when(() -> ZKMetadataProvider.getTableConfig(any(), anyString())).thenReturn(null); + zkMetadataProviderMockedStatic.when( + () -> ZKMetadataProvider.getSegmentZKMetadata(any(), anyString(), anyString())).thenReturn(segmentZKMetadata); + ZkHelixPropertyStore propertyStore = mock(ZkHelixPropertyStore.class); + HelixManager helixManager = mock(HelixManager.class); + when(helixManager.getHelixPropertyStore()).thenReturn(propertyStore); + + TableRebalancer.PeerDownloadTableDataLossRiskAssessor dataLossRiskAssessor = + new TableRebalancer.PeerDownloadTableDataLossRiskAssessor(REALTIME_TABLE_NAME, tableConfig, helixManager, + _helixResourceManager.getRealtimeSegmentManager()); + Pair dataLossRiskResult = dataLossRiskAssessor.assessDataLossRisk(segmentName); + assertFalse(dataLossRiskResult.getLeft()); + assertNull(dataLossRiskResult.getRight()); + } + + try (MockedStatic zkMetadataProviderMockedStatic = mockStatic(ZKMetadataProvider.class)) { + // Create segment with empty download URL + SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segmentName); + segmentZKMetadata.setStatus(CommonConstants.Segment.Realtime.Status.DONE); + segmentZKMetadata.setDownloadUrl(""); + + zkMetadataProviderMockedStatic.when(() -> ZKMetadataProvider.getTableConfig(any(), anyString())).thenReturn(null); + zkMetadataProviderMockedStatic.when( + () -> ZKMetadataProvider.getSegmentZKMetadata(any(), anyString(), anyString())).thenReturn(segmentZKMetadata); + ZkHelixPropertyStore propertyStore = mock(ZkHelixPropertyStore.class); + HelixManager helixManager = mock(HelixManager.class); + when(helixManager.getHelixPropertyStore()).thenReturn(propertyStore); + + TableRebalancer.PeerDownloadTableDataLossRiskAssessor dataLossRiskAssessor = + new TableRebalancer.PeerDownloadTableDataLossRiskAssessor(REALTIME_TABLE_NAME, tableConfig, helixManager, + _helixResourceManager.getRealtimeSegmentManager()); + Pair dataLossRiskResult = dataLossRiskAssessor.assessDataLossRisk(segmentName); + assertTrue(dataLossRiskResult.getLeft()); + assertTrue(dataLossRiskResult.getRight().contains("as the download URL is empty")); + } + + // Enable dedup on the table, this should return the same results with it disabled + DedupConfig dedupConfig = new DedupConfig(); + dedupConfig.setDedupEnabled(true); + tableConfig.setDedupConfig(dedupConfig); + + try (MockedStatic zkMetadataProviderMockedStatic = mockStatic(ZKMetadataProvider.class)) { + // Create segment with non-empty download URL + SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segmentName); + segmentZKMetadata.setStatus(CommonConstants.Segment.Realtime.Status.DONE); + segmentZKMetadata.setDownloadUrl("nonEmptyDownloadURL"); + + zkMetadataProviderMockedStatic.when(() -> ZKMetadataProvider.getTableConfig(any(), anyString())).thenReturn(null); + zkMetadataProviderMockedStatic.when( + () -> ZKMetadataProvider.getSegmentZKMetadata(any(), anyString(), anyString())).thenReturn(segmentZKMetadata); + ZkHelixPropertyStore propertyStore = mock(ZkHelixPropertyStore.class); + HelixManager helixManager = mock(HelixManager.class); + when(helixManager.getHelixPropertyStore()).thenReturn(propertyStore); + + TableRebalancer.PeerDownloadTableDataLossRiskAssessor dataLossRiskAssessor = + new TableRebalancer.PeerDownloadTableDataLossRiskAssessor(REALTIME_TABLE_NAME, tableConfig, helixManager, + _helixResourceManager.getRealtimeSegmentManager()); + Pair dataLossRiskResult = dataLossRiskAssessor.assessDataLossRisk(segmentName); + assertFalse(dataLossRiskResult.getLeft()); + assertNull(dataLossRiskResult.getRight()); + } + + try (MockedStatic zkMetadataProviderMockedStatic = mockStatic(ZKMetadataProvider.class)) { + // Create segment with empty download URL + SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segmentName); + segmentZKMetadata.setStatus(CommonConstants.Segment.Realtime.Status.DONE); + segmentZKMetadata.setDownloadUrl(""); + + zkMetadataProviderMockedStatic.when(() -> ZKMetadataProvider.getTableConfig(any(), anyString())).thenReturn(null); + zkMetadataProviderMockedStatic.when( + () -> ZKMetadataProvider.getSegmentZKMetadata(any(), anyString(), anyString())).thenReturn(segmentZKMetadata); + ZkHelixPropertyStore propertyStore = mock(ZkHelixPropertyStore.class); + HelixManager helixManager = mock(HelixManager.class); + when(helixManager.getHelixPropertyStore()).thenReturn(propertyStore); + + TableRebalancer.PeerDownloadTableDataLossRiskAssessor dataLossRiskAssessor = + new TableRebalancer.PeerDownloadTableDataLossRiskAssessor(REALTIME_TABLE_NAME, tableConfig, helixManager, + _helixResourceManager.getRealtimeSegmentManager()); + Pair dataLossRiskResult = dataLossRiskAssessor.assessDataLossRisk(segmentName); + assertTrue(dataLossRiskResult.getLeft()); + assertTrue(dataLossRiskResult.getRight().contains("as the download URL is empty")); + } + + // Enable upsert in PARTIAL mode on the table, this should return the same results with it disabled + tableConfig.setDedupConfig(null); + UpsertConfig upsertConfig = new UpsertConfig(); + upsertConfig.setMode(UpsertConfig.Mode.PARTIAL); + tableConfig.setUpsertConfig(upsertConfig); + + try (MockedStatic zkMetadataProviderMockedStatic = mockStatic(ZKMetadataProvider.class)) { + // Create segment with non-empty download URL + SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segmentName); + segmentZKMetadata.setStatus(CommonConstants.Segment.Realtime.Status.DONE); + segmentZKMetadata.setDownloadUrl("nonEmptyDownloadURL"); + + zkMetadataProviderMockedStatic.when(() -> ZKMetadataProvider.getTableConfig(any(), anyString())).thenReturn(null); + zkMetadataProviderMockedStatic.when( + () -> ZKMetadataProvider.getSegmentZKMetadata(any(), anyString(), anyString())).thenReturn(segmentZKMetadata); + ZkHelixPropertyStore propertyStore = mock(ZkHelixPropertyStore.class); + HelixManager helixManager = mock(HelixManager.class); + when(helixManager.getHelixPropertyStore()).thenReturn(propertyStore); + + TableRebalancer.PeerDownloadTableDataLossRiskAssessor dataLossRiskAssessor = + new TableRebalancer.PeerDownloadTableDataLossRiskAssessor(REALTIME_TABLE_NAME, tableConfig, helixManager, + _helixResourceManager.getRealtimeSegmentManager()); + Pair dataLossRiskResult = dataLossRiskAssessor.assessDataLossRisk(segmentName); + assertFalse(dataLossRiskResult.getLeft()); + assertNull(dataLossRiskResult.getRight()); + } + + try (MockedStatic zkMetadataProviderMockedStatic = mockStatic(ZKMetadataProvider.class)) { + // Create segment with empty download URL + SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segmentName); + segmentZKMetadata.setStatus(CommonConstants.Segment.Realtime.Status.DONE); + segmentZKMetadata.setDownloadUrl(""); + + zkMetadataProviderMockedStatic.when(() -> ZKMetadataProvider.getTableConfig(any(), anyString())).thenReturn(null); + zkMetadataProviderMockedStatic.when( + () -> ZKMetadataProvider.getSegmentZKMetadata(any(), anyString(), anyString())).thenReturn(segmentZKMetadata); + ZkHelixPropertyStore propertyStore = mock(ZkHelixPropertyStore.class); + HelixManager helixManager = mock(HelixManager.class); + when(helixManager.getHelixPropertyStore()).thenReturn(propertyStore); + + TableRebalancer.PeerDownloadTableDataLossRiskAssessor dataLossRiskAssessor = + new TableRebalancer.PeerDownloadTableDataLossRiskAssessor(REALTIME_TABLE_NAME, tableConfig, helixManager, + _helixResourceManager.getRealtimeSegmentManager()); + Pair dataLossRiskResult = dataLossRiskAssessor.assessDataLossRisk(segmentName); + assertTrue(dataLossRiskResult.getLeft()); + assertTrue(dataLossRiskResult.getRight().contains("as the download URL is empty")); + } + + // Enable pauseless, disable upsert and dedup, results should be the same as without pauseless enabled + tableConfig.setDedupConfig(null); + tableConfig.setUpsertConfig(null); + IngestionConfig ingestionConfig = new IngestionConfig(); + StreamIngestionConfig streamIngestionConfig = new StreamIngestionConfig( + Collections.singletonList(FakeStreamConfigUtils.getDefaultLowLevelStreamConfigs().getStreamConfigsMap())); + streamIngestionConfig.setPauselessConsumptionEnabled(true); + ingestionConfig.setStreamIngestionConfig(streamIngestionConfig); + tableConfig.setIngestionConfig(ingestionConfig); + + try (MockedStatic zkMetadataProviderMockedStatic = mockStatic(ZKMetadataProvider.class)) { + // Create segment with non-empty download URL + SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segmentName); + segmentZKMetadata.setStatus(CommonConstants.Segment.Realtime.Status.DONE); + segmentZKMetadata.setDownloadUrl("nonEmptyDownloadURL"); + + zkMetadataProviderMockedStatic.when(() -> ZKMetadataProvider.getTableConfig(any(), anyString())).thenReturn(null); + zkMetadataProviderMockedStatic.when( + () -> ZKMetadataProvider.getSegmentZKMetadata(any(), anyString(), anyString())).thenReturn(segmentZKMetadata); + ZkHelixPropertyStore propertyStore = mock(ZkHelixPropertyStore.class); + HelixManager helixManager = mock(HelixManager.class); + when(helixManager.getHelixPropertyStore()).thenReturn(propertyStore); + + TableRebalancer.PeerDownloadTableDataLossRiskAssessor dataLossRiskAssessor = + new TableRebalancer.PeerDownloadTableDataLossRiskAssessor(REALTIME_TABLE_NAME, tableConfig, helixManager, + _helixResourceManager.getRealtimeSegmentManager()); + Pair dataLossRiskResult = dataLossRiskAssessor.assessDataLossRisk(segmentName); + assertFalse(dataLossRiskResult.getLeft()); + assertNull(dataLossRiskResult.getRight()); + } + + try (MockedStatic zkMetadataProviderMockedStatic = mockStatic(ZKMetadataProvider.class)) { + // Create segment with empty download URL + SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segmentName); + segmentZKMetadata.setStatus(CommonConstants.Segment.Realtime.Status.DONE); + segmentZKMetadata.setDownloadUrl(""); + + zkMetadataProviderMockedStatic.when(() -> ZKMetadataProvider.getTableConfig(any(), anyString())).thenReturn(null); + zkMetadataProviderMockedStatic.when( + () -> ZKMetadataProvider.getSegmentZKMetadata(any(), anyString(), anyString())).thenReturn(segmentZKMetadata); + ZkHelixPropertyStore propertyStore = mock(ZkHelixPropertyStore.class); + HelixManager helixManager = mock(HelixManager.class); + when(helixManager.getHelixPropertyStore()).thenReturn(propertyStore); + + TableRebalancer.PeerDownloadTableDataLossRiskAssessor dataLossRiskAssessor = + new TableRebalancer.PeerDownloadTableDataLossRiskAssessor(REALTIME_TABLE_NAME, tableConfig, helixManager, + _helixResourceManager.getRealtimeSegmentManager()); + Pair dataLossRiskResult = dataLossRiskAssessor.assessDataLossRisk(segmentName); + assertTrue(dataLossRiskResult.getLeft()); + assertTrue(dataLossRiskResult.getRight().contains("as the download URL is empty")); + } + } + + @Test + public void testDataLossRiskAssessorPeerDownloadEnabledConsumingSegment() { + String segmentName = "randomSegmentName"; + TableConfig tableConfig = + new TableConfigBuilder(TableType.REALTIME).setTableName(RAW_TABLE_NAME).setNumReplicas(NUM_REPLICAS).build(); + tableConfig.getValidationConfig().setPeerSegmentDownloadScheme("http"); + + try (MockedStatic zkMetadataProviderMockedStatic = mockStatic(ZKMetadataProvider.class)) { + // Create IN_PROGRESS segment with empty download URL. Download URL should not exist for segment that's not yet + // completed + SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segmentName); + segmentZKMetadata.setStatus(CommonConstants.Segment.Realtime.Status.IN_PROGRESS); + segmentZKMetadata.setDownloadUrl(""); + + zkMetadataProviderMockedStatic.when(() -> ZKMetadataProvider.getTableConfig(any(), anyString())).thenReturn(null); + zkMetadataProviderMockedStatic.when( + () -> ZKMetadataProvider.getSegmentZKMetadata(any(), anyString(), anyString())).thenReturn(segmentZKMetadata); + ZkHelixPropertyStore propertyStore = mock(ZkHelixPropertyStore.class); + HelixManager helixManager = mock(HelixManager.class); + when(helixManager.getHelixPropertyStore()).thenReturn(propertyStore); + + TableRebalancer.PeerDownloadTableDataLossRiskAssessor dataLossRiskAssessor = + new TableRebalancer.PeerDownloadTableDataLossRiskAssessor(REALTIME_TABLE_NAME, tableConfig, helixManager, + _helixResourceManager.getRealtimeSegmentManager()); + Pair dataLossRiskResult = dataLossRiskAssessor.assessDataLossRisk(segmentName); + assertFalse(dataLossRiskResult.getLeft()); + assertNull(dataLossRiskResult.getRight()); + } + + // Enable dedup on the table, this should return the same results with it disabled + DedupConfig dedupConfig = new DedupConfig(); + dedupConfig.setDedupEnabled(true); + tableConfig.setDedupConfig(dedupConfig); + + try (MockedStatic zkMetadataProviderMockedStatic = mockStatic(ZKMetadataProvider.class)) { + // Create IN_PROGRESS segment with empty download URL. Download URL should not exist for segment that's not yet + // completed + SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segmentName); + segmentZKMetadata.setStatus(CommonConstants.Segment.Realtime.Status.IN_PROGRESS); + segmentZKMetadata.setDownloadUrl(""); + + zkMetadataProviderMockedStatic.when(() -> ZKMetadataProvider.getTableConfig(any(), anyString())).thenReturn(null); + zkMetadataProviderMockedStatic.when( + () -> ZKMetadataProvider.getSegmentZKMetadata(any(), anyString(), anyString())).thenReturn(segmentZKMetadata); + ZkHelixPropertyStore propertyStore = mock(ZkHelixPropertyStore.class); + HelixManager helixManager = mock(HelixManager.class); + when(helixManager.getHelixPropertyStore()).thenReturn(propertyStore); + + TableRebalancer.PeerDownloadTableDataLossRiskAssessor dataLossRiskAssessor = + new TableRebalancer.PeerDownloadTableDataLossRiskAssessor(REALTIME_TABLE_NAME, tableConfig, helixManager, + _helixResourceManager.getRealtimeSegmentManager()); + Pair dataLossRiskResult = dataLossRiskAssessor.assessDataLossRisk(segmentName); + assertFalse(dataLossRiskResult.getLeft()); + assertNull(dataLossRiskResult.getRight()); + } + + // Enable upsert in PARTIAL mode on the table, this should return the same results with it disabled + tableConfig.setDedupConfig(null); + UpsertConfig upsertConfig = new UpsertConfig(); + upsertConfig.setMode(UpsertConfig.Mode.PARTIAL); + tableConfig.setUpsertConfig(upsertConfig); + + try (MockedStatic zkMetadataProviderMockedStatic = mockStatic(ZKMetadataProvider.class)) { + // Create IN_PROGRESS segment with empty download URL. Download URL should not exist for segment that's not yet + // completed + SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segmentName); + segmentZKMetadata.setStatus(CommonConstants.Segment.Realtime.Status.IN_PROGRESS); + segmentZKMetadata.setDownloadUrl(""); + + zkMetadataProviderMockedStatic.when(() -> ZKMetadataProvider.getTableConfig(any(), anyString())).thenReturn(null); + zkMetadataProviderMockedStatic.when( + () -> ZKMetadataProvider.getSegmentZKMetadata(any(), anyString(), anyString())).thenReturn(segmentZKMetadata); + ZkHelixPropertyStore propertyStore = mock(ZkHelixPropertyStore.class); + HelixManager helixManager = mock(HelixManager.class); + when(helixManager.getHelixPropertyStore()).thenReturn(propertyStore); + + TableRebalancer.PeerDownloadTableDataLossRiskAssessor dataLossRiskAssessor = + new TableRebalancer.PeerDownloadTableDataLossRiskAssessor(REALTIME_TABLE_NAME, tableConfig, helixManager, + _helixResourceManager.getRealtimeSegmentManager()); + Pair dataLossRiskResult = dataLossRiskAssessor.assessDataLossRisk(segmentName); + assertFalse(dataLossRiskResult.getLeft()); + assertNull(dataLossRiskResult.getRight()); + } + + // Enable pauseless, disable dedup and upsert, this should return the same results as with it disabled + tableConfig.setDedupConfig(null); + tableConfig.setUpsertConfig(null); + IngestionConfig ingestionConfig = new IngestionConfig(); + StreamIngestionConfig streamIngestionConfig = new StreamIngestionConfig( + Collections.singletonList(FakeStreamConfigUtils.getDefaultLowLevelStreamConfigs().getStreamConfigsMap())); + streamIngestionConfig.setPauselessConsumptionEnabled(true); + ingestionConfig.setStreamIngestionConfig(streamIngestionConfig); + tableConfig.setIngestionConfig(ingestionConfig); + + try (MockedStatic zkMetadataProviderMockedStatic = mockStatic(ZKMetadataProvider.class)) { + // Create IN_PROGRESS segment with empty download URL. Download URL should not exist for segment that's not yet + // completed + SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segmentName); + segmentZKMetadata.setStatus(CommonConstants.Segment.Realtime.Status.IN_PROGRESS); + segmentZKMetadata.setDownloadUrl(""); + + zkMetadataProviderMockedStatic.when(() -> ZKMetadataProvider.getTableConfig(any(), anyString())).thenReturn(null); + zkMetadataProviderMockedStatic.when( + () -> ZKMetadataProvider.getSegmentZKMetadata(any(), anyString(), anyString())).thenReturn(segmentZKMetadata); + ZkHelixPropertyStore propertyStore = mock(ZkHelixPropertyStore.class); + HelixManager helixManager = mock(HelixManager.class); + when(helixManager.getHelixPropertyStore()).thenReturn(propertyStore); + + TableRebalancer.PeerDownloadTableDataLossRiskAssessor dataLossRiskAssessor = + new TableRebalancer.PeerDownloadTableDataLossRiskAssessor(REALTIME_TABLE_NAME, tableConfig, helixManager, + _helixResourceManager.getRealtimeSegmentManager()); + Pair dataLossRiskResult = dataLossRiskAssessor.assessDataLossRisk(segmentName); + assertFalse(dataLossRiskResult.getLeft()); + assertNull(dataLossRiskResult.getRight()); + } + } + + @Test + public void testDataLossRiskAssessorPeerDownloadEnabledPauselessEnabledCommittingSegment() { + String segmentName = "randomSegmentName"; + TableConfig tableConfig = + new TableConfigBuilder(TableType.REALTIME).setTableName(RAW_TABLE_NAME).setNumReplicas(NUM_REPLICAS).build(); + tableConfig.getValidationConfig().setPeerSegmentDownloadScheme("https"); + + // Enable pauseless, this by itself (without dedup / upsert) should return false + // No need to test non-pauseless tables with COMMITTING state as this is only applicable to pauseless tables + IngestionConfig ingestionConfig = new IngestionConfig(); + StreamIngestionConfig streamIngestionConfig = new StreamIngestionConfig( + Collections.singletonList(FakeStreamConfigUtils.getDefaultLowLevelStreamConfigs().getStreamConfigsMap())); + streamIngestionConfig.setPauselessConsumptionEnabled(true); + ingestionConfig.setStreamIngestionConfig(streamIngestionConfig); + tableConfig.setIngestionConfig(ingestionConfig); + + try (MockedStatic zkMetadataProviderMockedStatic = mockStatic(ZKMetadataProvider.class)) { + // Create segment in COMMITTING state + SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segmentName); + segmentZKMetadata.setStatus(CommonConstants.Segment.Realtime.Status.COMMITTING); + segmentZKMetadata.setDownloadUrl(""); + + zkMetadataProviderMockedStatic.when(() -> ZKMetadataProvider.getTableConfig(any(), anyString())).thenReturn(null); + zkMetadataProviderMockedStatic.when( + () -> ZKMetadataProvider.getSegmentZKMetadata(any(), anyString(), anyString())).thenReturn(segmentZKMetadata); + ZkHelixPropertyStore propertyStore = mock(ZkHelixPropertyStore.class); + HelixManager helixManager = mock(HelixManager.class); + when(helixManager.getHelixPropertyStore()).thenReturn(propertyStore); + + TableRebalancer.PeerDownloadTableDataLossRiskAssessor dataLossRiskAssessor = + new TableRebalancer.PeerDownloadTableDataLossRiskAssessor(REALTIME_TABLE_NAME, tableConfig, helixManager, + _helixResourceManager.getRealtimeSegmentManager()); + Pair dataLossRiskResult = dataLossRiskAssessor.assessDataLossRisk(segmentName); + assertFalse(dataLossRiskResult.getLeft()); + assertNull(dataLossRiskResult.getRight()); + } + + // Enable dedup on the table, this should return true + DedupConfig dedupConfig = new DedupConfig(); + dedupConfig.setDedupEnabled(true); + tableConfig.setDedupConfig(dedupConfig); + + try (MockedStatic zkMetadataProviderMockedStatic = mockStatic(ZKMetadataProvider.class)) { + // Create segment with empty download URL + SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segmentName); + segmentZKMetadata.setStatus(CommonConstants.Segment.Realtime.Status.COMMITTING); + segmentZKMetadata.setDownloadUrl(""); + + zkMetadataProviderMockedStatic.when(() -> ZKMetadataProvider.getTableConfig(any(), anyString())).thenReturn(null); + zkMetadataProviderMockedStatic.when( + () -> ZKMetadataProvider.getSegmentZKMetadata(any(), anyString(), anyString())).thenReturn(segmentZKMetadata); + ZkHelixPropertyStore propertyStore = mock(ZkHelixPropertyStore.class); + HelixManager helixManager = mock(HelixManager.class); + when(helixManager.getHelixPropertyStore()).thenReturn(propertyStore); + + TableRebalancer.PeerDownloadTableDataLossRiskAssessor dataLossRiskAssessor = + new TableRebalancer.PeerDownloadTableDataLossRiskAssessor(REALTIME_TABLE_NAME, tableConfig, helixManager, + _helixResourceManager.getRealtimeSegmentManager()); + Pair dataLossRiskResult = dataLossRiskAssessor.assessDataLossRisk(segmentName); + assertTrue(dataLossRiskResult.getLeft()); + assertTrue(dataLossRiskResult.getRight().contains("as it is in COMMITING state")); + } + + // Enable upsert on the table in PARTIAL mode, this should return true + tableConfig.setDedupConfig(null); + UpsertConfig upsertConfig = new UpsertConfig(); + upsertConfig.setMode(UpsertConfig.Mode.PARTIAL); + tableConfig.setUpsertConfig(upsertConfig); + + try (MockedStatic zkMetadataProviderMockedStatic = mockStatic(ZKMetadataProvider.class)) { + // Create segment with empty download URL + SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segmentName); + segmentZKMetadata.setStatus(CommonConstants.Segment.Realtime.Status.COMMITTING); + segmentZKMetadata.setDownloadUrl(""); + + zkMetadataProviderMockedStatic.when(() -> ZKMetadataProvider.getTableConfig(any(), anyString())).thenReturn(null); + zkMetadataProviderMockedStatic.when( + () -> ZKMetadataProvider.getSegmentZKMetadata(any(), anyString(), anyString())).thenReturn(segmentZKMetadata); + ZkHelixPropertyStore propertyStore = mock(ZkHelixPropertyStore.class); + HelixManager helixManager = mock(HelixManager.class); + when(helixManager.getHelixPropertyStore()).thenReturn(propertyStore); + + TableRebalancer.PeerDownloadTableDataLossRiskAssessor dataLossRiskAssessor = + new TableRebalancer.PeerDownloadTableDataLossRiskAssessor(REALTIME_TABLE_NAME, tableConfig, helixManager, + _helixResourceManager.getRealtimeSegmentManager()); + Pair dataLossRiskResult = dataLossRiskAssessor.assessDataLossRisk(segmentName); + assertTrue(dataLossRiskResult.getLeft()); + assertTrue(dataLossRiskResult.getRight().contains("as it is in COMMITING state")); + } + + // Enable upsert on the table in FULL mode, this should return false + upsertConfig = new UpsertConfig(); + upsertConfig.setMode(UpsertConfig.Mode.FULL); + tableConfig.setUpsertConfig(upsertConfig); + + try (MockedStatic zkMetadataProviderMockedStatic = mockStatic(ZKMetadataProvider.class)) { + // Create segment with empty download URL + SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segmentName); + segmentZKMetadata.setStatus(CommonConstants.Segment.Realtime.Status.COMMITTING); + segmentZKMetadata.setDownloadUrl(""); + + zkMetadataProviderMockedStatic.when(() -> ZKMetadataProvider.getTableConfig(any(), anyString())).thenReturn(null); + zkMetadataProviderMockedStatic.when( + () -> ZKMetadataProvider.getSegmentZKMetadata(any(), anyString(), anyString())).thenReturn(segmentZKMetadata); + ZkHelixPropertyStore propertyStore = mock(ZkHelixPropertyStore.class); + HelixManager helixManager = mock(HelixManager.class); + when(helixManager.getHelixPropertyStore()).thenReturn(propertyStore); + + TableRebalancer.PeerDownloadTableDataLossRiskAssessor dataLossRiskAssessor = + new TableRebalancer.PeerDownloadTableDataLossRiskAssessor(REALTIME_TABLE_NAME, tableConfig, helixManager, + _helixResourceManager.getRealtimeSegmentManager()); + Pair dataLossRiskResult = dataLossRiskAssessor.assessDataLossRisk(segmentName); + assertFalse(dataLossRiskResult.getLeft()); + assertNull(dataLossRiskResult.getRight()); + } + } +} diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerClusterStatelessTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerClusterStatelessTest.java index 2f96d5cf35ec..dbb87f8f8933 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerClusterStatelessTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerClusterStatelessTest.java @@ -54,12 +54,11 @@ import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType; import org.apache.pinot.spi.config.table.assignment.InstanceReplicaGroupPartitionConfig; import org.apache.pinot.spi.config.table.assignment.InstanceTagPoolConfig; -import org.apache.pinot.spi.config.table.ingestion.IngestionConfig; -import org.apache.pinot.spi.config.table.ingestion.StreamIngestionConfig; import org.apache.pinot.spi.config.tenant.Tenant; import org.apache.pinot.spi.config.tenant.TenantRole; import org.apache.pinot.spi.stream.LongMsgOffset; import org.apache.pinot.spi.stream.StreamPartitionMsgOffset; +import org.apache.pinot.spi.utils.CommonConstants; import org.apache.pinot.spi.utils.Enablement; import org.apache.pinot.spi.utils.builder.TableConfigBuilder; import org.apache.pinot.spi.utils.builder.TableNameBuilder; @@ -670,6 +669,90 @@ public void testRebalance() } } + @Test + public void testRebalancePeerDownloadDataLoss() + throws Exception { + for (int batchSizePerServer : Arrays.asList(RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, 1, 2)) { + addFakeServerInstanceToAutoJoinHelixCluster(SERVER_INSTANCE_ID_PREFIX, true); + + ExecutorService executorService = Executors.newFixedThreadPool(10); + DefaultRebalancePreChecker preChecker = new DefaultRebalancePreChecker(); + preChecker.init(_helixResourceManager, executorService, 1); + TableRebalancer tableRebalancer = + new TableRebalancer(_helixManager, null, null, preChecker, _tableSizeReader, null); + TableConfig tableConfig = + new TableConfigBuilder(TableType.REALTIME).setTableName(RAW_TABLE_NAME) + .setNumReplicas(1) + .setStreamConfigs(FakeStreamConfigUtils.getDefaultLowLevelStreamConfigs().getStreamConfigsMap()) + .build(); + // Create the table + addDummySchema(RAW_TABLE_NAME); + _helixResourceManager.addTable(tableConfig); + + // Add the segments with peer download uri (i.e. simulate segments without deep store uri) + int numSegments = 10; + for (int i = 0; i < numSegments; i++) { + _helixResourceManager.addNewSegment(REALTIME_TABLE_NAME, + SegmentMetadataMockUtils.mockSegmentMetadata(RAW_TABLE_NAME, SEGMENT_NAME_PREFIX + i), + CommonConstants.Segment.METADATA_URI_FOR_PEER_DOWNLOAD); + } + + // Rebalance should return NO_OP status + RebalanceConfig rebalanceConfig = new RebalanceConfig(); + rebalanceConfig.setBatchSizePerServer(batchSizePerServer); + RebalanceResult rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.NO_OP); + + // Add 1 more servers + addFakeServerInstanceToAutoJoinHelixCluster(SERVER_INSTANCE_ID_PREFIX + 1, true); + + // Enable peer-download for the table and validate that rebalance with different parameters + // The table has segments without deep store uri, so it should fail with peer download data loss protection + + // Case 1: downtime=true, allowPeerDownloadDataLoss=false (should fail) + tableConfig.getValidationConfig().setPeerSegmentDownloadScheme("http"); + rebalanceConfig = new RebalanceConfig(); + rebalanceConfig.setDowntime(true); + rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.FAILED); + // to make sure the failure is due to peer download data loss protection (the description is too long, only check + // the keyword here) + assertTrue(rebalanceResult.getDescription().toLowerCase().contains("peer-download")); + + // Case 2: downtime=true, allowPeerDownloadDataLoss=true (should succeed) + rebalanceConfig.setDowntime(true); + rebalanceConfig.setAllowPeerDownloadDataLoss(true); + rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); + + // Case 3: downtime=false, minAvailableReplicas=0, allowPeerDownloadDataLoss=false (should fail) + // Add 1 more servers + addFakeServerInstanceToAutoJoinHelixCluster(SERVER_INSTANCE_ID_PREFIX + 2, true); + + rebalanceConfig.setDowntime(false); + rebalanceConfig.setAllowPeerDownloadDataLoss(false); + rebalanceConfig.setMinAvailableReplicas(0); + rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.FAILED); + // to make sure the failure is due to peer download data loss protection + assertTrue(rebalanceResult.getDescription().toLowerCase().contains("peer-download")); + + // Case 4: downtime=false, minAvailableReplicas=0, allowPeerDownloadDataLoss=true (should succeed) + rebalanceConfig.setAllowPeerDownloadDataLoss(true); + rebalanceResult = tableRebalancer.rebalance(tableConfig, rebalanceConfig, null); + // notice that in real world scenario, the rebalance will hang because servers cannot find segments to download + // in this stateless test, servers are mocked to always succeed any state transition so we expect a DONE here + assertEquals(rebalanceResult.getStatus(), RebalanceResult.Status.DONE); + + _helixResourceManager.deleteRealtimeTable(RAW_TABLE_NAME); + + stopAndDropFakeInstance(SERVER_INSTANCE_ID_PREFIX); + stopAndDropFakeInstance(SERVER_INSTANCE_ID_PREFIX + 1); + stopAndDropFakeInstance(SERVER_INSTANCE_ID_PREFIX + 2); + executorService.shutdown(); + } + } + @Test(timeOut = 60000) public void testRebalanceStrictReplicaGroup() throws Exception { @@ -1207,13 +1290,8 @@ public void testRebalancePreCheckerRebalanceConfig() assertEquals(preCheckerResult.getPreCheckStatus(), RebalancePreCheckerResult.PreCheckStatus.PASS); assertEquals(preCheckerResult.getMessage(), "All rebalance parameters look good"); - // trigger pauseless table rebalance warning - IngestionConfig ingestionConfig = new IngestionConfig(); - StreamIngestionConfig streamIngestionConfig = new StreamIngestionConfig( - Collections.singletonList(FakeStreamConfigUtils.getDefaultLowLevelStreamConfigs().getStreamConfigsMap())); - streamIngestionConfig.setPauselessConsumptionEnabled(true); - ingestionConfig.setStreamIngestionConfig(streamIngestionConfig); - newTableConfig.setIngestionConfig(ingestionConfig); + // trigger peer-download enabled table rebalance warning + newTableConfig.getValidationConfig().setPeerSegmentDownloadScheme("http"); rebalanceConfig.setDowntime(true); rebalanceResult = tableRebalancer.rebalance(newTableConfig, rebalanceConfig, null); @@ -1221,19 +1299,19 @@ public void testRebalancePreCheckerRebalanceConfig() assertNotNull(preCheckerResult); assertEquals(preCheckerResult.getPreCheckStatus(), RebalancePreCheckerResult.PreCheckStatus.WARN); assertEquals(preCheckerResult.getMessage(), - "Replication of the table is 1, which is not recommended for pauseless tables as it may cause data loss " - + "during rebalance"); + "Replication of the table is 1, which is not recommended for peer-download enabled tables as it may " + + "cause data loss during rebalance"); newTableConfig = new TableConfigBuilder(TableType.REALTIME).setTableName(RAW_TABLE_NAME).setNumReplicas(3).build(); - newTableConfig.setIngestionConfig(ingestionConfig); + newTableConfig.getValidationConfig().setPeerSegmentDownloadScheme("https"); rebalanceResult = tableRebalancer.rebalance(newTableConfig, rebalanceConfig, null); preCheckerResult = rebalanceResult.getPreChecksResult().get(DefaultRebalancePreChecker.REBALANCE_CONFIG_OPTIONS); assertNotNull(preCheckerResult); assertEquals(preCheckerResult.getPreCheckStatus(), RebalancePreCheckerResult.PreCheckStatus.WARN); assertEquals(preCheckerResult.getMessage(), - "Number of replicas (3) is greater than 1, downtime is not recommended.\nDowntime or minAvailableReplicas=0 " - + "for pauseless tables may cause data loss during rebalance"); + "Number of replicas (3) is greater than 1, downtime is not recommended.\nDowntime or minAvailableReplicas<=0 " + + "for peer-download enabled tables may cause data loss during rebalance"); rebalanceConfig.setDowntime(false); rebalanceConfig.setMinAvailableReplicas(-3); @@ -1242,7 +1320,7 @@ public void testRebalancePreCheckerRebalanceConfig() assertNotNull(preCheckerResult); assertEquals(preCheckerResult.getPreCheckStatus(), RebalancePreCheckerResult.PreCheckStatus.WARN); assertEquals(preCheckerResult.getMessage(), - "Downtime or minAvailableReplicas=0 for pauseless tables may cause data loss during rebalance"); + "Downtime or minAvailableReplicas<=0 for peer-download enabled tables may cause data loss during rebalance"); rebalanceConfig.setDowntime(false); rebalanceConfig.setMinAvailableReplicas(0); @@ -1251,7 +1329,7 @@ public void testRebalancePreCheckerRebalanceConfig() assertNotNull(preCheckerResult); assertEquals(preCheckerResult.getPreCheckStatus(), RebalancePreCheckerResult.PreCheckStatus.WARN); assertEquals(preCheckerResult.getMessage(), - "Downtime or minAvailableReplicas=0 for pauseless tables may cause data loss during rebalance"); + "Downtime or minAvailableReplicas<=0 for peer-download enabled tables may cause data loss during rebalance"); // test pass rebalanceConfig.setMinAvailableReplicas(1); diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerTest.java index f59972c9a529..635ba1657a7e 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerTest.java @@ -38,6 +38,7 @@ import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; import static org.testng.Assert.assertNotEquals; +import static org.testng.Assert.assertThrows; import static org.testng.Assert.assertTrue; import static org.testng.Assert.fail; @@ -50,6 +51,11 @@ public class TableRebalancerTest { return name == null ? -1 : name.getPartitionGroupId(); }; + private static final TableRebalancer.DataLossRiskAssessor DEFAULT_DATA_LOSS_RISK_ASSESSOR = + new TableRebalancer.NoOpRiskAssessor(); + private static final TableRebalancer.DataLossRiskAssessor ALWAYS_TRUE_DATA_LOSS_RISK_ASSESSOR = + segmentName -> Pair.of(true, ""); + @Test public void testDowntimeMode() { // With common instance, first assignment should be the same as target assignment @@ -662,7 +668,8 @@ public void testAssignment() { for (boolean enableStrictReplicaGroup : Arrays.asList(false, true)) { Map> nextAssignment = TableRebalancer.getNextAssignment(currentAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER); + RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER, + DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertEquals(nextAssignment, targetAssignment); } @@ -749,7 +756,8 @@ public void testAssignment() { for (boolean enableStrictReplicaGroup : Arrays.asList(false, true)) { Map> nextAssignment = TableRebalancer.getNextAssignment(currentAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER); + RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER, + DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertEquals(nextAssignment.get("segment1").keySet(), new TreeSet<>(Arrays.asList("host1", "host2", "host4"))); assertEquals(nextAssignment.get("segment2").keySet(), new TreeSet<>(Arrays.asList("host2", "host4", "host5"))); assertEquals(nextAssignment.get("segment3").keySet(), new TreeSet<>(Arrays.asList("host1", "host2", "host4"))); @@ -757,7 +765,8 @@ public void testAssignment() { nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER); + RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER, + DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertEquals(nextAssignment, targetAssignment); } @@ -811,7 +820,8 @@ public void testAssignment() { // Next assignment with 2 minimum available replicas without strict replica-group should reach the target assignment Map> nextAssignment = TableRebalancer.getNextAssignment(currentAssignment, targetAssignment, 2, false, false, - RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER); + RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER, + DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertEquals(nextAssignment, targetAssignment); // Next assignment with 2 minimum available replicas with strict replica-group should finish in 2 steps: @@ -844,13 +854,15 @@ public void testAssignment() { // // The second assignment should reach the target assignment nextAssignment = TableRebalancer.getNextAssignment(currentAssignment, targetAssignment, 2, true, false, - RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER); + RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER, + DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertEquals(nextAssignment.get("segment1").keySet(), new TreeSet<>(Arrays.asList("host1", "host3", "host4"))); assertEquals(nextAssignment.get("segment2").keySet(), new TreeSet<>(Arrays.asList("host2", "host3", "host4"))); assertEquals(nextAssignment.get("segment3").keySet(), new TreeSet<>(Arrays.asList("host1", "host3", "host4"))); assertEquals(nextAssignment.get("segment4").keySet(), new TreeSet<>(Arrays.asList("host2", "host3", "host4"))); nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, true, false, - RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER); + RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER, + DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertEquals(nextAssignment, targetAssignment); } @@ -968,7 +980,8 @@ public void testAssignmentWithLowDiskMode() { for (boolean enableStrictReplicaGroup : Arrays.asList(false, true)) { Map> nextAssignment = TableRebalancer.getNextAssignment(currentAssignment, targetAssignment, 2, enableStrictReplicaGroup, true, - RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER); + RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER, + DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertEquals(nextAssignment.get("segment1").keySet(), new TreeSet<>(Arrays.asList("host1", "host3"))); assertEquals(nextAssignment.get("segment2").keySet(), new TreeSet<>(Arrays.asList("host2", "host4"))); assertEquals(nextAssignment.get("segment3").keySet(), new TreeSet<>(Arrays.asList("host1", "host3"))); @@ -976,7 +989,8 @@ public void testAssignmentWithLowDiskMode() { nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, enableStrictReplicaGroup, true, - RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER); + RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER, + DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertEquals(nextAssignment, targetAssignment); } @@ -1104,7 +1118,8 @@ public void testAssignmentWithLowDiskMode() { for (boolean enableStrictReplicaGroup : Arrays.asList(false, true)) { Map> nextAssignment = TableRebalancer.getNextAssignment(currentAssignment, targetAssignment, 2, enableStrictReplicaGroup, true, - RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER); + RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER, + DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertEquals(nextAssignment.get("segment1").keySet(), new TreeSet<>(Arrays.asList("host1", "host2"))); assertEquals(nextAssignment.get("segment2").keySet(), new TreeSet<>(Arrays.asList("host2", "host4"))); assertEquals(nextAssignment.get("segment3").keySet(), new TreeSet<>(Arrays.asList("host1", "host2"))); @@ -1112,7 +1127,8 @@ public void testAssignmentWithLowDiskMode() { nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, enableStrictReplicaGroup, true, - RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER); + RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER, + DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertEquals(nextAssignment.get("segment1").keySet(), new TreeSet<>(Arrays.asList("host1", "host2", "host4"))); assertEquals(nextAssignment.get("segment2").keySet(), new TreeSet<>(Arrays.asList("host2", "host4", "host5"))); assertEquals(nextAssignment.get("segment3").keySet(), new TreeSet<>(Arrays.asList("host1", "host2", "host4"))); @@ -1120,7 +1136,8 @@ public void testAssignmentWithLowDiskMode() { nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, enableStrictReplicaGroup, true, - RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER); + RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER, + DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertEquals(nextAssignment.get("segment1").keySet(), new TreeSet<>(Arrays.asList("host2", "host4"))); assertEquals(nextAssignment.get("segment2").keySet(), new TreeSet<>(Arrays.asList("host4", "host5"))); assertEquals(nextAssignment.get("segment3").keySet(), new TreeSet<>(Arrays.asList("host2", "host4"))); @@ -1128,7 +1145,8 @@ public void testAssignmentWithLowDiskMode() { nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, enableStrictReplicaGroup, true, - RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER); + RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER, + DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertEquals(nextAssignment, targetAssignment); } @@ -1205,14 +1223,16 @@ public void testAssignmentWithLowDiskMode() { // The second assignment should reach the target assignment Map> nextAssignment = TableRebalancer.getNextAssignment(currentAssignment, targetAssignment, 2, false, true, - RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER); + RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER, + DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertEquals(nextAssignment.get("segment1").keySet(), new TreeSet<>(Arrays.asList("host1", "host3"))); assertEquals(nextAssignment.get("segment2").keySet(), new TreeSet<>(Arrays.asList("host3", "host4"))); assertEquals(nextAssignment.get("segment3").keySet(), new TreeSet<>(Arrays.asList("host1", "host3"))); assertEquals(nextAssignment.get("segment4").keySet(), new TreeSet<>(Arrays.asList("host3", "host4"))); nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, false, true, - RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER); + RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER, + DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertEquals(nextAssignment, targetAssignment); // Next assignment with 2 minimum available replicas with strict replica-group should finish in 3 steps: @@ -1264,21 +1284,24 @@ public void testAssignmentWithLowDiskMode() { // // The third assignment should reach the target assignment nextAssignment = TableRebalancer.getNextAssignment(currentAssignment, targetAssignment, 2, true, true, - RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER); + RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER, + DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertEquals(nextAssignment.get("segment1").keySet(), new TreeSet<>(Arrays.asList("host1", "host3"))); assertEquals(nextAssignment.get("segment2").keySet(), new TreeSet<>(Arrays.asList("host3", "host4"))); assertEquals(nextAssignment.get("segment3").keySet(), new TreeSet<>(Arrays.asList("host1", "host3"))); assertEquals(nextAssignment.get("segment4").keySet(), new TreeSet<>(Arrays.asList("host3", "host4"))); nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, true, true, - RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER); + RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER, + DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertEquals(nextAssignment.get("segment1").keySet(), new TreeSet<>(Arrays.asList("host1", "host3", "host4"))); assertEquals(nextAssignment.get("segment2").keySet(), new TreeSet<>(Arrays.asList("host3", "host4"))); assertEquals(nextAssignment.get("segment3").keySet(), new TreeSet<>(Arrays.asList("host1", "host3", "host4"))); assertEquals(nextAssignment.get("segment4").keySet(), new TreeSet<>(Arrays.asList("host3", "host4"))); nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, true, true, - RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER); + RebalanceConfig.DISABLE_BATCH_SIZE_PER_SERVER, new Object2IntOpenHashMap<>(), DUMMY_PARTITION_FETCHER, + DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertEquals(nextAssignment, targetAssignment); } @@ -1573,7 +1596,7 @@ public void testAssignmentWithServerBatching() { Object2IntOpenHashMap segmentToPartitionIdMap = new Object2IntOpenHashMap<>(); Map> nextAssignment = TableRebalancer.getNextAssignment(currentAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - 1, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER); + 1, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertNotEquals(nextAssignment, targetAssignment); assertEquals(nextAssignment.get("segment__1__0__98347869999L").keySet(), new TreeSet<>(Arrays.asList("host1", "host3", "host5"))); @@ -1585,7 +1608,7 @@ public void testAssignmentWithServerBatching() { new TreeSet<>(Arrays.asList("host2", "host3", "host4"))); nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - 1, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER); + 1, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertEquals(nextAssignment, targetAssignment); } @@ -1595,7 +1618,7 @@ public void testAssignmentWithServerBatching() { Object2IntOpenHashMap segmentToPartitionIdMap = new Object2IntOpenHashMap<>(); Map> nextAssignment = TableRebalancer.getNextAssignment(currentAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - 2, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER); + 2, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertEquals(nextAssignment, targetAssignment); } @@ -1753,7 +1776,7 @@ public void testAssignmentWithServerBatching() { for (boolean enableStrictReplicaGroup : Arrays.asList(false, true)) { Map> nextAssignment = TableRebalancer.getNextAssignment(currentAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - 1, new Object2IntOpenHashMap<>(), SIMPLE_PARTITION_FETCHER); + 1, new Object2IntOpenHashMap<>(), SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertEquals(nextAssignment.get("segment__1__0__98347869999L").keySet(), new TreeSet<>(Arrays.asList("host1", "host2", "host4"))); assertEquals(nextAssignment.get("segment__2__0__98347869999L").keySet(), @@ -1765,7 +1788,7 @@ public void testAssignmentWithServerBatching() { nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - 1, new Object2IntOpenHashMap<>(), SIMPLE_PARTITION_FETCHER); + 1, new Object2IntOpenHashMap<>(), SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); if (!enableStrictReplicaGroup) { assertEquals(nextAssignment.get("segment__1__0__98347869999L").keySet(), new TreeSet<>(Arrays.asList("host2", "host4", "host6"))); @@ -1777,7 +1800,7 @@ public void testAssignmentWithServerBatching() { new TreeSet<>(Arrays.asList("host2", "host4", "host5"))); nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - 1, new Object2IntOpenHashMap<>(), SIMPLE_PARTITION_FETCHER); + 1, new Object2IntOpenHashMap<>(), SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); } else { assertEquals(nextAssignment.get("segment__1__0__98347869999L").keySet(), new TreeSet<>(Arrays.asList("host1", "host2", "host4"))); @@ -1789,7 +1812,7 @@ public void testAssignmentWithServerBatching() { new TreeSet<>(Arrays.asList("host2", "host4", "host5"))); nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - 1, new Object2IntOpenHashMap<>(), SIMPLE_PARTITION_FETCHER); + 1, new Object2IntOpenHashMap<>(), SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertEquals(nextAssignment.get("segment__1__0__98347869999L").keySet(), new TreeSet<>(Arrays.asList("host1", "host2", "host4"))); assertEquals(nextAssignment.get("segment__2__0__98347869999L").keySet(), @@ -1800,7 +1823,7 @@ public void testAssignmentWithServerBatching() { new TreeSet<>(Arrays.asList("host1", "host4", "host5"))); nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - 1, new Object2IntOpenHashMap<>(), SIMPLE_PARTITION_FETCHER); + 1, new Object2IntOpenHashMap<>(), SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); } assertEquals(nextAssignment, targetAssignment); } @@ -1856,7 +1879,7 @@ public void testAssignmentWithServerBatching() { // in 1 step using batchSizePerServer = 2 Map> nextAssignment = TableRebalancer.getNextAssignment(currentAssignment, targetAssignment, 2, false, false, - 2, new Object2IntOpenHashMap<>(), SIMPLE_PARTITION_FETCHER); + 2, new Object2IntOpenHashMap<>(), SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertEquals(nextAssignment, targetAssignment); // Next assignment with 2 minimum available replicas with strict replica-group should finish in 2 steps even with @@ -1890,7 +1913,7 @@ public void testAssignmentWithServerBatching() { // // The second assignment should reach the target assignment nextAssignment = TableRebalancer.getNextAssignment(currentAssignment, targetAssignment, 2, true, false, - 2, new Object2IntOpenHashMap<>(), SIMPLE_PARTITION_FETCHER); + 2, new Object2IntOpenHashMap<>(), SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertEquals(nextAssignment.get("segment__1__0__98347869999L").keySet(), new TreeSet<>(Arrays.asList("host1", "host2", "host3"))); assertEquals(nextAssignment.get("segment__2__0__98347869999L").keySet(), @@ -1900,7 +1923,7 @@ public void testAssignmentWithServerBatching() { assertEquals(nextAssignment.get("segment__4__0__98347869999L").keySet(), new TreeSet<>(Arrays.asList("host1", "host3", "host4"))); nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, true, false, - 2, new Object2IntOpenHashMap<>(), SIMPLE_PARTITION_FETCHER); + 2, new Object2IntOpenHashMap<>(), SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertEquals(nextAssignment, targetAssignment); // Try assignment with overlapping partitions across segments, especially for strict replica group based. @@ -1998,7 +2021,7 @@ public void testAssignmentWithServerBatching() { Object2IntOpenHashMap segmentToPartitionIdMap = new Object2IntOpenHashMap<>(); nextAssignment = TableRebalancer.getNextAssignment(currentAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - 1, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER); + 1, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); if (!enableStrictReplicaGroup) { // Nothing should change, since we don't select based on partitions for non-strict replica groups assertNotEquals(nextAssignment, targetAssignment); @@ -2012,7 +2035,7 @@ public void testAssignmentWithServerBatching() { new TreeSet<>(Arrays.asList("host2", "host3", "host4"))); nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - 1, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER); + 1, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); } assertEquals(nextAssignment, targetAssignment); } @@ -2077,7 +2100,7 @@ public void testAssignmentWithServerBatching() { Object2IntOpenHashMap segmentToPartitionIdMap = new Object2IntOpenHashMap<>(); nextAssignment = TableRebalancer.getNextAssignment(currentAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - 1, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER); + 1, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); if (!enableStrictReplicaGroup) { // Nothing should change, since we don't select based on partitions for non-strict replica groups assertNotEquals(nextAssignment, targetAssignment); @@ -2091,7 +2114,7 @@ public void testAssignmentWithServerBatching() { new TreeSet<>(Arrays.asList("host2", "host3", "host4"))); nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - 1, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER); + 1, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertNotEquals(nextAssignment, targetAssignment); assertEquals(nextAssignment.get("segment__1__0__98347869999L").keySet(), new TreeSet<>(Arrays.asList("host2", "host4", "host6"))); @@ -2103,7 +2126,7 @@ public void testAssignmentWithServerBatching() { new TreeSet<>(Arrays.asList("host2", "host4", "host5"))); nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - 1, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER); + 1, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); } else { assertNotEquals(nextAssignment, targetAssignment); assertEquals(nextAssignment.get("segment__1__0__98347869999L").keySet(), @@ -2116,7 +2139,7 @@ public void testAssignmentWithServerBatching() { new TreeSet<>(Arrays.asList("host2", "host4", "host5"))); nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - 1, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER); + 1, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); } assertEquals(nextAssignment, targetAssignment); } @@ -2245,7 +2268,7 @@ public void testAssignmentWithServerBatching() { Object2IntOpenHashMap segmentToPartitionIdMap = new Object2IntOpenHashMap<>(); nextAssignment = TableRebalancer.getNextAssignment(currentAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - 1, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER); + 1, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); if (!enableStrictReplicaGroup) { assertNotEquals(nextAssignment, targetAssignment); assertEquals(nextAssignment.get("segment__1__0__98347869999L").keySet(), @@ -2262,7 +2285,7 @@ public void testAssignmentWithServerBatching() { new TreeSet<>(Arrays.asList("host1", "host2", "host3"))); nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - 1, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER); + 1, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertNotEquals(nextAssignment, targetAssignment); assertEquals(nextAssignment.get("segment__1__0__98347869999L").keySet(), new TreeSet<>(Arrays.asList("host1", "host3", "host5"))); @@ -2278,7 +2301,7 @@ public void testAssignmentWithServerBatching() { new TreeSet<>(Arrays.asList("host1", "host2", "host3"))); nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - 1, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER); + 1, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertNotEquals(nextAssignment, targetAssignment); assertEquals(nextAssignment.get("segment__1__0__98347869999L").keySet(), new TreeSet<>(Arrays.asList("host1", "host3", "host5"))); @@ -2294,7 +2317,7 @@ public void testAssignmentWithServerBatching() { new TreeSet<>(Arrays.asList("host1", "host2", "host3"))); nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - 1, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER); + 1, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); } else { assertNotEquals(nextAssignment, targetAssignment); assertEquals(nextAssignment.get("segment__1__0__98347869999L").keySet(), @@ -2311,7 +2334,7 @@ public void testAssignmentWithServerBatching() { new TreeSet<>(Arrays.asList("host1", "host2", "host3"))); nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - 1, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER); + 1, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); } assertEquals(nextAssignment, targetAssignment); } @@ -2484,7 +2507,7 @@ public void testAssignmentWithServerBatching() { Object2IntOpenHashMap segmentToPartitionIdMap = new Object2IntOpenHashMap<>(); nextAssignment = TableRebalancer.getNextAssignment(currentAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - 2, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER); + 2, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); if (!enableStrictReplicaGroup) { assertNotEquals(nextAssignment, targetAssignment); assertEquals(nextAssignment.get("segment__1__0__98347869999L").keySet(), @@ -2507,7 +2530,7 @@ public void testAssignmentWithServerBatching() { new TreeSet<>(Arrays.asList("host1", "host3", "host6"))); nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - 2, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER); + 2, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertNotEquals(nextAssignment, targetAssignment); assertEquals(nextAssignment.get("segment__1__0__98347869999L").keySet(), new TreeSet<>(Arrays.asList("host1", "host3", "host5"))); @@ -2529,7 +2552,7 @@ public void testAssignmentWithServerBatching() { new TreeSet<>(Arrays.asList("host1", "host3", "host6"))); nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - 2, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER); + 2, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); } else { assertNotEquals(nextAssignment, targetAssignment); assertEquals(nextAssignment.get("segment__1__0__98347869999L").keySet(), @@ -2552,7 +2575,7 @@ public void testAssignmentWithServerBatching() { new TreeSet<>(Arrays.asList("host1", "host3", "host5"))); nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - 2, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER); + 2, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); assertNotEquals(nextAssignment, targetAssignment); assertEquals(nextAssignment.get("segment__1__0__98347869999L").keySet(), new TreeSet<>(Arrays.asList("host1", "host3", "host5"))); @@ -2574,7 +2597,7 @@ public void testAssignmentWithServerBatching() { new TreeSet<>(Arrays.asList("host1", "host3", "host5"))); nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - 2, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER); + 2, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); } assertEquals(nextAssignment, targetAssignment); } @@ -2585,7 +2608,7 @@ public void testAssignmentWithServerBatching() { Object2IntOpenHashMap segmentToPartitionIdMap = new Object2IntOpenHashMap<>(); nextAssignment = TableRebalancer.getNextAssignment(currentAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - 5, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER); + 5, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); if (!enableStrictReplicaGroup) { assertNotEquals(nextAssignment, targetAssignment); assertEquals(nextAssignment.get("segment__1__0__98347869999L").keySet(), @@ -2609,7 +2632,7 @@ public void testAssignmentWithServerBatching() { new TreeSet<>(Arrays.asList("host1", "host3", "host6"))); nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - 2, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER); + 2, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); } else { // This would have completed in a single step for batchSizePerServer = 5 if we did not have an additional check // to only add segments that exceed the batchSizePerServer if no prior segments were already assigned to a @@ -2635,9 +2658,51 @@ public void testAssignmentWithServerBatching() { new TreeSet<>(Arrays.asList("host1", "host3", "host5"))); nextAssignment = TableRebalancer.getNextAssignment(nextAssignment, targetAssignment, 2, enableStrictReplicaGroup, false, - 2, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER); + 2, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER, DEFAULT_DATA_LOSS_RISK_ASSESSOR); } assertEquals(nextAssignment, targetAssignment); } } + + @Test + public void testAssignmentWithDataLossAssessorReturnTrue() { + // Test that exception is thrown if the DataLossRiskAssessor returns true + Map> currentAssignment = new TreeMap<>(); + currentAssignment.put("segment__1__0__98347869999L", + SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host1", "host2", "host3"), ONLINE)); + currentAssignment.put("segment__2__0__98347869999L", + SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host2", "host3", "host4"), ONLINE)); + currentAssignment.put("segment__3__0__98347869999L", + SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host1", "host2", "host3"), ONLINE)); + currentAssignment.put("segment__4__0__98347869999L", + SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host2", "host3", "host4"), ONLINE)); + + Map> targetAssignment = new TreeMap<>(); + targetAssignment.put("segment__1__0__98347869999L", + SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host1", "host3", "host5"), ONLINE)); + targetAssignment.put("segment__2__0__98347869999L", + SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host2", "host4", "host6"), ONLINE)); + targetAssignment.put("segment__3__0__98347869999L", + SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host1", "host3", "host5"), ONLINE)); + targetAssignment.put("segment__4__0__98347869999L", + SegmentAssignmentUtils.getInstanceStateMap(Arrays.asList("host2", "host4", "host6"), ONLINE)); + + Map numSegmentsToOffloadMap = + TableRebalancer.getNumSegmentsToOffloadMap(currentAssignment, targetAssignment); + assertEquals(numSegmentsToOffloadMap.size(), 6); + assertEquals((int) numSegmentsToOffloadMap.get("host1"), 0); + assertEquals((int) numSegmentsToOffloadMap.get("host2"), 2); + assertEquals((int) numSegmentsToOffloadMap.get("host3"), 2); + assertEquals((int) numSegmentsToOffloadMap.get("host4"), 0); + assertEquals((int) numSegmentsToOffloadMap.get("host5"), -2); + assertEquals((int) numSegmentsToOffloadMap.get("host6"), -2); + + // The DataLossRiskAssessor returns true, so we expect an exception to be thrown + for (boolean enableStrictReplicaGroup : Arrays.asList(false, true)) { + Object2IntOpenHashMap segmentToPartitionIdMap = new Object2IntOpenHashMap<>(); + assertThrows(IllegalStateException.class, () -> TableRebalancer.getNextAssignment(currentAssignment, + targetAssignment, 2, enableStrictReplicaGroup, false, 2, segmentToPartitionIdMap, SIMPLE_PARTITION_FETCHER, + ALWAYS_TRUE_DATA_LOSS_RISK_ASSESSOR)); + } + } } From 94adc1a6d5feda3b2b9158d0f0663d76524bfea6 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 7 Aug 2025 12:30:58 -0700 Subject: [PATCH 099/167] Bump net.openhft:chronicle-bom from 2.27ea65 to 2.27ea66 (#16546) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 9d12ed3a194c..c02b52d28e77 100644 --- a/pom.xml +++ b/pom.xml @@ -296,7 +296,7 @@ 2.2.0 5.0.4 5.5.1 - 2.27ea65 + 2.27ea66 2.0.6.1 3.9.11 2.2.0 From befd8dfc777108a67047888e452a1b114845dca4 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 7 Aug 2025 12:34:04 -0700 Subject: [PATCH 100/167] Bump org.jline:jline from 3.30.4 to 3.30.5 (#16548) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index c02b52d28e77..36aba1d11d87 100644 --- a/pom.xml +++ b/pom.xml @@ -254,7 +254,7 @@ 2.1.0 - 3.30.4 + 3.30.5 2.0.1 1.5.4 10.4.1 From d51026471729ca40be08e53b58f7aa4ffcb01a87 Mon Sep 17 00:00:00 2001 From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com> Date: Thu, 7 Aug 2025 12:34:36 -0700 Subject: [PATCH 101/167] Bump org.webjars:swagger-ui from 5.27.0 to 5.27.1 (#16549) --- pom.xml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pom.xml b/pom.xml index 36aba1d11d87..e714b40fd95e 100644 --- a/pom.xml +++ b/pom.xml @@ -151,7 +151,7 @@ 2.47 2.6.1 1.6.16 - 5.27.0 + 5.27.1 3.4.1 2.9.0 2.6.0 From 8d061eea14f7d717e75a831864b2cb8467348ee8 Mon Sep 17 00:00:00 2001 From: Saptarshi Sengupta <94242536+saptarshi1996@users.noreply.github.com> Date: Fri, 8 Aug 2025 01:07:46 +0530 Subject: [PATCH 102/167] rebalance ui flow fixes (#16545) --- .../RebalanceServer/RebalanceResponse.tsx | 24 ++- .../Operations/RebalanceServerStatusOp.tsx | 57 +++++- .../Operations/RebalanceServerTableOp.tsx | 169 ++++++++++++------ 3 files changed, 181 insertions(+), 69 deletions(-) diff --git a/pinot-controller/src/main/resources/app/components/Homepage/Operations/RebalanceServer/RebalanceResponse.tsx b/pinot-controller/src/main/resources/app/components/Homepage/Operations/RebalanceServer/RebalanceResponse.tsx index d8d09a60d6f6..67658974100b 100644 --- a/pinot-controller/src/main/resources/app/components/Homepage/Operations/RebalanceServer/RebalanceResponse.tsx +++ b/pinot-controller/src/main/resources/app/components/Homepage/Operations/RebalanceServer/RebalanceResponse.tsx @@ -54,7 +54,7 @@ const CopyJobIdToClipboardButton = ({ jobId }: { jobId: string }) => { ); }; -export const RebalanceResponse = ({ response }) => { +export const RebalanceResponse = ({ response, onJobIdClick }) => { const responseSectionsToShow = [ { name: 'Segment Assignment', @@ -81,9 +81,25 @@ export const RebalanceResponse = ({ response }) => { Job Id - - {response.jobId} - + {response.status === 'IN_PROGRESS' && onJobIdClick ? ( + onJobIdClick(response.jobId)} + title="Click to view job details" + > + {response.jobId} + + ) : ( + + {response.jobId} + + )} diff --git a/pinot-controller/src/main/resources/app/components/Homepage/Operations/RebalanceServerStatusOp.tsx b/pinot-controller/src/main/resources/app/components/Homepage/Operations/RebalanceServerStatusOp.tsx index 9d61cf6f5ae5..4a07b4a0ba98 100644 --- a/pinot-controller/src/main/resources/app/components/Homepage/Operations/RebalanceServerStatusOp.tsx +++ b/pinot-controller/src/main/resources/app/components/Homepage/Operations/RebalanceServerStatusOp.tsx @@ -35,10 +35,11 @@ import {RebalanceTableSegmentJob} from "Models"; type RebalanceServerStatusOpProps = { tableName: string; hideModal: () => void; + initialJobId?: string; }; export const RebalanceServerStatusOp = ( - { tableName, hideModal } : RebalanceServerStatusOpProps + { tableName, hideModal, initialJobId } : RebalanceServerStatusOpProps ) => { const { currentTimezone } = useTimezone(); const [rebalanceServerJobs, setRebalanceServerJobs] = React.useState([]) @@ -47,7 +48,7 @@ export const RebalanceServerStatusOp = ( const [rebalanceProgressStats, setRebalanceProgressStats] = useState<{}>({}); const [loading, setLoading] = useState(false); - useEffect(() => { + const fetchRebalanceJobs = () => { setLoading(true); PinotMethodUtils .fetchRebalanceTableJobs(tableName) @@ -55,8 +56,35 @@ export const RebalanceServerStatusOp = ( setRebalanceServerJobs(jobs) }) .finally(() => setLoading(false)); + }; + + useEffect(() => { + fetchRebalanceJobs(); }, []); + // Set initial job selection when jobs are loaded and initialJobId is provided + useEffect(() => { + if (initialJobId && rebalanceServerJobs.length > 0 && !jobSelected) { + const jobExists = rebalanceServerJobs.find(job => job.jobId === initialJobId); + if (jobExists) { + setJobSelected(initialJobId); + } + } + }, [initialJobId, rebalanceServerJobs, jobSelected]); + + const RefreshAction = () => { + return ( + + ); + }; + const BackAction = () => { return ( ); } +const RebalanceAction = ({ handleOnRun, disabled }: { handleOnRun: () => void, disabled?: boolean }) => { + return ( + + ); +} + const BackAction = ({ onClick }: { onClick: () => void }) => { return (